Python class

Weights

`Weights`

class max.graph.weights.Weights(*args, **kwargs)

source

Bases: Protocol

Protocol for managing and accessing model weights hierarchically.

The Weights protocol provides a convenient interface for loading and organizing neural network weights. It supports hierarchical naming through attribute and index access, making it easy to work with complex model architectures.

Weights in MAX are tensors backed by external memory (buffers or memory-mapped files) that remain separate from the compiled graph.

from max.graph import Graph
from max.dtype import DType

graph = Graph("my_model")
weights = graph.weights()

# Allocate weights with hierarchical naming
attn_weight = weights.transformer.layers[0].attention.weight.allocate(
    dtype=DType.float32,
    shape=(768, 768)
)
# Creates weight named "transformer.layers.0.attention.weight"

# Check if a weight exists before allocating
if weights.transformer.layers[0].mlp.weight.exists():
    mlp_weight = weights.transformer.layers[0].mlp.weight.allocate(
        dtype=DType.float16,
        shape=(768, 3072)
    )

`allocate()`

allocate(dtype=None, shape=None, quantization_encoding=None, device=cpu:0)

source

Creates a Weight object for this tensor.

# Allocate a weight with specific configuration
weight = weights.model.layers[0].weight.allocate(
    dtype=DType.float16,  # Convert to half precision
    shape=(768, 768),
    device=DeviceRef.GPU(0)  # Place on first GPU
)

# Add to graph
with graph:
    weight_tensor = graph.add_weight(weight)

Parameters:

dtype (DType | None) – Data type for the weight. If None, uses the original dtype.
shape (Iterable[int | str | Dim | integer[Any] | TypedAttr] | None) – Shape of the weight tensor. If None, uses the original shape.
quantization_encoding (QuantizationEncoding | None) – Quantization scheme to apply (for example, Q4_K, Q8_0).
device (DeviceRef) – Target device for the weight (CPU or GPU).

Returns:

A Weight object that can be added to a graph using graph.add_weight().

Return type:

Weight

`allocated_weights`

property allocated_weights: dict[str, DLPackArray]

source

Returns all previously allocated weights.

This only includes weights that were explicitly allocated using Weights.allocate(), not all available weights.

Returns:: A dictionary mapping weight names to their numpy arrays for all weights that have been allocated through this interface.

`data()`

data()

source

Returns weight data with metadata.

weight_data = weights.model.embeddings.weight.data()
print(f"Shape: {weight_data.shape}")
print(f"Dtype: {weight_data.dtype}")

# Convert to different dtype
fp16_data = weight_data.astype(DType.float16)

Returns:: A WeightData object containing the tensor data along with metadata like name, dtype, shape, and quantization encoding.
Raises:: KeyError – If no weight exists at the current hierarchical name.
Return type:: WeightData

`exists()`

exists()

source

Checks if a weight with this exact name exists.

if weights.model.classifier.weight.exists():
    classifier = weights.model.classifier.weight.allocate(...)
else:
    print("Classifier weight not found")

Returns:: True if a weight with the current hierarchical name exists in the loaded weights, False otherwise.
Return type:: bool

`items()`

items()

source

Iterates through all weights that start with the current prefix.

# Iterate through all weights in a specific layer
for name, weight in weights.transformer.layers[0].items():
    print(f"Found weight: {name}")

Yields:: Tuples of (name, weight_accessor) for each weight under the current prefix. The name is relative to the current prefix.
Parameters:: self (_Self)
Return type:: Iterator[tuple[str, _Self]]

`name`

property name: str

source

The current weight name or prefix.

Returns:: The hierarchical name built from attribute and index access. For example, if accessed as weights.model.layers[0], returns model.layers.0.

Weights​

allocate()​

allocated_weights​

data()​

exists()​

items()​

name​