gm.nn.IntWrapper

class gemma.gm.nn.IntWrapper(

*,

model: flax.linen.module.Module,

dtype: numpy.dtype = <class 'jax.numpy.int4'>,

parent: flax.linen.module.Module | flax.core.scope.Scope | flax.linen.module._Sentinel | None = <flax.linen.module._Sentinel object>,

name: str | None = None,

)

Bases: flax.linen.module.Module

围绕 Gemma 模型的外包层，用于启用 int4 推理。

被包裹的模型将替换所有 nn.Dense、nn.Einsum 等层为其 int4 版本。有关更多详细信息，请参阅 gemma.peft 文档。

model

要包裹的模型。

类型:

flax.linen.module.Module

model: flax.linen.module.Module

dtype

alias of jax.numpy.int4

name: str | None = None

parent: flax.linen.module.Module | flax.core.scope.Scope | flax.linen.module._Sentinel | None = None

scope: Scope | None = None