For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /get-started.md).

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.

import json
import struct
import tempfile
from pathlib import Path

import numpy as np
from max.dtype import DType
from max.graph import DeviceRef
from max.graph.weights import load_weights

def write_safetensors(path, tensors):
    header, buffers, offset = {}, [], 0
    for name, arr in tensors.items():
        arr = np.ascontiguousarray(arr)
        header[name] = {
            "dtype": "F32",
            "shape": list(arr.shape),
            "data_offsets": [offset, offset + arr.nbytes],
        }
        buffers.append(arr.tobytes())
        offset += arr.nbytes
    blob = json.dumps(header).encode()
    with open(path, "wb") as f:
        f.write(struct.pack("<Q", len(blob)))
        f.write(blob)
        for b in buffers:
            f.write(b)

with tempfile.TemporaryDirectory() as tmp:
    path = Path(tmp) / "model.safetensors"
    write_safetensors(
        path,
        {
            "transformer.layers.0.attention.weight": np.ones(
                (8, 8), dtype=np.float32
            ),
        },
    )
    weights = load_weights([path])

    attn_weight = weights.transformer.layers[0].attention.weight.allocate(
        dtype=DType.float32,
        device=DeviceRef.CPU(),
    )
    # Creates weight named "transformer.layers.0.attention.weight".

allocate()

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

source

Creates a Weight object for this tensor.

import json
import struct
import tempfile
from pathlib import Path

import numpy as np
from max.dtype import DType
from max.graph import DeviceRef, Graph
from max.graph.weights import load_weights

def write_safetensors(path, tensors):
    header, buffers, offset = {}, [], 0
    for name, arr in tensors.items():
        arr = np.ascontiguousarray(arr)
        header[name] = {
            "dtype": "F32",
            "shape": list(arr.shape),
            "data_offsets": [offset, offset + arr.nbytes],
        }
        buffers.append(arr.tobytes())
        offset += arr.nbytes
    blob = json.dumps(header).encode()
    with open(path, "wb") as f:
        f.write(struct.pack("<Q", len(blob)))
        f.write(blob)
        for b in buffers:
            f.write(b)

with tempfile.TemporaryDirectory() as tmp:
    path = Path(tmp) / "model.safetensors"
    write_safetensors(
        path,
        {
            "model.layers.0.weight": np.ones(
                (8, 8), dtype=np.float32
            )
        },
    )
    weights = load_weights([path])

    weight = weights.model.layers[0].weight.allocate(
        dtype=DType.float32,
        shape=(8, 8),
        device=DeviceRef.CPU(),
    )

    with Graph("allocate_example", input_types=[]) as 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 | 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.

import json
import struct
import tempfile
from pathlib import Path

import numpy as np
from max.dtype import DType
from max.graph.weights import load_weights

def write_safetensors(path, tensors):
    header, buffers, offset = {}, [], 0
    for name, arr in tensors.items():
        arr = np.ascontiguousarray(arr)
        header[name] = {
            "dtype": "F32",
            "shape": list(arr.shape),
            "data_offsets": [offset, offset + arr.nbytes],
        }
        buffers.append(arr.tobytes())
        offset += arr.nbytes
    blob = json.dumps(header).encode()
    with open(path, "wb") as f:
        f.write(struct.pack("<Q", len(blob)))
        f.write(blob)
        for b in buffers:
            f.write(b)

with tempfile.TemporaryDirectory() as tmp:
    path = Path(tmp) / "model.safetensors"
    write_safetensors(
        path,
        {
            "model.embeddings.weight": np.ones(
                (4, 4), dtype=np.float32
            )
        },
    )
    weights = load_weights([path])

    weight_data = weights.model.embeddings.weight.data()
    # weight_data.shape and weight_data.dtype hold the metadata.

    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.

import json
import struct
import tempfile
from pathlib import Path

import numpy as np
from max.dtype import DType
from max.graph import DeviceRef
from max.graph.weights import load_weights

def write_safetensors(path, tensors):
    header, buffers, offset = {}, [], 0
    for name, arr in tensors.items():
        arr = np.ascontiguousarray(arr)
        header[name] = {
            "dtype": "F32",
            "shape": list(arr.shape),
            "data_offsets": [offset, offset + arr.nbytes],
        }
        buffers.append(arr.tobytes())
        offset += arr.nbytes
    blob = json.dumps(header).encode()
    with open(path, "wb") as f:
        f.write(struct.pack("<Q", len(blob)))
        f.write(blob)
        for b in buffers:
            f.write(b)

with tempfile.TemporaryDirectory() as tmp:
    path = Path(tmp) / "model.safetensors"
    write_safetensors(
        path, {"model.head.weight": np.ones((4, 4), dtype=np.float32)}
    )
    weights = load_weights([path])

    if weights.model.head.weight.exists():
        head = weights.model.head.weight.allocate(
            dtype=DType.float32, device=DeviceRef.CPU()
        )
    else:
        head = None

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.

import json
import struct
import tempfile
from pathlib import Path

import numpy as np
from max.graph.weights import load_weights

def write_safetensors(path, tensors):
    header, buffers, offset = {}, [], 0
    for name, arr in tensors.items():
        arr = np.ascontiguousarray(arr)
        header[name] = {
            "dtype": "F32",
            "shape": list(arr.shape),
            "data_offsets": [offset, offset + arr.nbytes],
        }
        buffers.append(arr.tobytes())
        offset += arr.nbytes
    blob = json.dumps(header).encode()
    with open(path, "wb") as f:
        f.write(struct.pack("<Q", len(blob)))
        f.write(blob)
        for b in buffers:
            f.write(b)

with tempfile.TemporaryDirectory() as tmp:
    path = Path(tmp) / "model.safetensors"
    write_safetensors(
        path,
        {
            "transformer.layers.0.q_proj.weight": np.ones(
                (4, 4), dtype=np.float32
            ),
            "transformer.layers.0.k_proj.weight": np.ones(
                (4, 4), dtype=np.float32
            ),
        },
    )
    weights = load_weights([path])

    found = [
        name
        for name, weight in weights.transformer.layers[0].items()
    ]

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.