Python class

Weight

`Weight`

class max.graph.Weight(*args, **kwargs)

source

Bases: TensorValue

Represents a value in a Graph that can be loaded at a later time.

Weights can be initialized outside of a Graph and are lazily-added to the parent graph when used. If there is no parent graph when a weight is used, an error will be raised.

Parameters:

name – The name of the weight.
dtype – The data type of the weight.
shape – The shape of the weight.
device – The device where the weight resides.
quantization_encoding – Optional quantization encoding for the weight.
align – Optional alignment requirement in bytes.
sharding_strategy – Optional sharding strategy for distributed execution.

Initializes a TensorValue from a tensor-like value.

Parameters:: value – The value to wrap. Can be an MLIR tensor value, another TensorValue, a Dim, or a Shape.

`align`

align: int | None

source

`device`

property device: DeviceRef

source

The device where the weight resides.

`dtype`

property dtype: DType

source

The data type of the weight.

`original_dtype_and_shape`

property original_dtype_and_shape: tuple[DType, Shape]

source

The original dtype and shape of this weight.

This property should be used to store the original weight’s dtype and shape the quantization encoding forces the weight to be loaded as uint8.

`quantization_encoding`

quantization_encoding: QuantizationEncoding | None

source

`shape`

property shape: Shape

source

The shape of the weight.

For sharded weights, returns the shape of the shard. Otherwise, returns the original weight shape.

`shard()`

shard(devices)

source

Creates sharded views of this Weight across multiple devices.

This Weight must have sharding_strategy defined. The shard objects returned are also Weight objects, but cannot be sharded further.

Parameters:: devices (Iterable[DeviceRef]) – Iterable of devices to place the shards on.
Returns:: List of sharded weights, one for each device.
Return type:: list[Weight]

`shard_idx`

shard_idx: int | None

source

`sharding_strategy`

property sharding_strategy: ShardingStrategy | None

source

Gets the weight sharding strategy.

Weight​

align​

device​

dtype​

original_dtype_and_shape​

quantization_encoding​

shape​

shard()​

shard_idx​

sharding_strategy​