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

SafetensorWeights

SafetensorWeights

class max.graph.weights.SafetensorWeights(filepaths, *, tensors=None, tensors_to_file_idx=None, prefix='', allocated=None, _st_weight_map=None, _st_file_handles=None)

source

Bases: Weights

Implementation for loading weights from safetensors files.

SafetensorWeights provides a secure and efficient way to load model weights from safetensors format files. Safetensors is designed by Hugging Face for safe serialization that prevents arbitrary code execution and supports memory-mapped loading for fast access.

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 SafetensorWeights

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
            ),
            "transformer.layers.0.attention.weight": np.ones(
                (4, 4), dtype=np.float32
            ),
        },
    )

    weights = SafetensorWeights([path])

    if weights.model.embeddings.weight.exists():
        embedding_weight = weights.model.embeddings.weight.allocate(
            dtype=DType.float32,
            device=DeviceRef.CPU(),
        )

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

Parameters:

• filepaths (Sequence[PathLike[str]])
• tensors (Set[str] | None)
• tensors_to_file_idx (Mapping[str, int] | None)
• prefix (str)
• allocated (dict[str, DLPackArray] | None)
• _st_weight_map (dict[str, Buffer])
• _st_file_handles (dict[PathLike[str], SafeTensor])

allocate()

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

source

Creates a Weight that can be added to a graph.

Parameters:

• dtype (DType | None)
• shape (Iterable[int | str | Dim | integer | TypedAttr] | None)
• quantization_encoding (QuantizationEncoding | None)
• device (DeviceRef)

Return type:

Weight

allocate_as_bytes()

allocate_as_bytes(dtype=None)

source

Creates a Weight that can be added to the graph with uint8 representation.

The last dimension is scaled by the number of bytes of the original dtype (for example, [512, 256] float32 becomes [512, 1024] uint8). Scalars are interpreted as shape [1].

Parameters:

dtype (DType | None)

Return type:

Weight

allocated_weights

property allocated_weights: dict[str, DLPackArray]

source

Gets the values of all weights that were allocated previously.

data()

data()

source

Loads and returns the weight data for this tensor.

Return type:

WeightData

exists()

exists()

source

Returns True if a tensor exists for the current name.

Return type:

bool

items()

items()

source

Iterates through all allocable weights that start with the prefix.

name

property name: str

source

The current weight name or prefix.