Python class

WeightScaleSpec

`WeightScaleSpec`

class max.nn.WeightScaleSpec(granularity, dtype, block_size=None)

source

Specifies how weights are scaled for scaled quantization.

Parameters:

block_size: tuple[int, int] | None = None

source

The tuple[int, int] of the block size for block-wise scaling.

dtype: DType

source

The DType of the weight scale factor(s).

granularity: ScaleGranularity

source

The ScaleGranularity of the weight scale factor application.

property is_block: bool

source

Whether the weight scale granularity is block-wise.

property is_colwise: bool

source

Whether the weight scale granularity is column-wise.

property is_rowwise: bool

source

Whether the weight scale granularity is row-wise.

property is_tensor: bool

source

Whether the weight scale granularity is per-tensor.