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 function

load_weights

`load_weights()`

max.graph.weights.load_weights(paths)

source

Loads neural network weights from checkpoint files.

Automatically detects checkpoint formats based on file extensions and returns the appropriate Weights implementation. Supported formats:

.safetensors (Safetensors)
.gguf (GGUF)

The following example shows how to load weights from a Safetensors file:

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.layers.23.mlp.gate_proj.weight": np.ones(
                (16, 8), dtype=np.float32
            )
        },
    )

    # load_weights also accepts multiple paths for sharded checkpoints.
    weights = load_weights([path])
    layer_weight = weights.model.layers[23].mlp.gate_proj.weight.allocate(
        dtype=DType.float32,
        shape=[16, 8],
        device=DeviceRef.CPU(),
    )

Parameters:: paths (list[Path]) – List of pathlib.Path objects pointing to checkpoint files. For multi-file checkpoints (for example, sharded Safetensors), provide all file paths in the list. For single-file checkpoints, provide a list with one path.
Return type:: Weights

load_weights()​

`load_weights()`