Skip to main content

Python class

Weight

Weight

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

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.

Initialize a Weight.

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.
  • _placeholder – Internal flag indicating if this is a placeholder weight.
  • _has_alias – Internal flag indicating if this weight has an alias.

align

align: int | None

device

property device: DeviceRef

The device where the weight resides.

dtype

property dtype: DType

The data type of the weight.

original_dtype_and_shape

property original_dtype_and_shape: tuple[DType, Shape]

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

shape

property shape: Shape

The shape of the weight.

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

shard()

shard(devices)

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

sharding_strategy

property sharding_strategy: ShardingStrategy | None

Gets the weight sharding strategy.

Was this page helpful?