Python module
weights
APIs for loading weights into a graph.
GGUFWeights
class max.graph.weights.GGUFWeights(source, tensors=None, prefix='', allocated=None)
Creates a GGUF weights reader.
-
Parameters:
-
- source (
Union[PathLike,gguf.GGUFReader]) – Path to a GGUF file or a GGUFReader object. - tensors – List of tensors in the GGUF checkpoint.
- prefix (
str) – Weight name or prefix. - allocated – Dictionary of allocated values.
- source (
allocate()
allocate(dtype=None, shape=None, quantization_encoding=None, device=cpu:0)
Creates and optionally validates a new Weight.
allocated_weights
property allocated_weights: dict[str, ndarray[Any, dtype[_ScalarType_co]]]
Gets the values of all weights that were allocated previously.
data()
data()
Returns data loaded from the weights at the current prefix.
-
Raises:
-
KeyError if the current prefix isn't present in the checkpoint. –
-
Return type:
exists()
exists()
Returns whether a weight with this exact name exists.
-
Return type:
items()
items()
Iterate through all allocable weights that start with the prefix.
name
property name: str
The current weight name or prefix.
raw_tensor()
raw_tensor()
Returns the numpy tensor corresponding to this weights object.
PytorchWeights
class max.graph.weights.PytorchWeights(filepath, tensor_infos=None, prefix='', allocated=None)
allocate()
allocate(dtype=None, shape=None, quantization_encoding=None, device=cpu:0)
Creates and optionally validates a new Weight.
allocated_weights
property allocated_weights: dict[str, ndarray[Any, dtype[_ScalarType_co]]]
Gets the values of all weights that were allocated previously.
data()
data()
-
Return type:
dtype
property dtype: DType
The current weight dtype, if this weight exists.
exists()
exists()
-
Return type:
items()
items()
Iterate through all allocable weights that start with the prefix.
name
property name: str
The current weight name or prefix.
quantization_encoding
property quantization_encoding: QuantizationEncoding | None
The current weight quantization encoding, if this weight exists.
raw_tensor()
raw_tensor()
Returns the tensor corresponding to this weights object.
shape
property shape: Shape
The current weight shape, if this weight exists.
RandomWeights
class max.graph.weights.RandomWeights(_allocated=<factory>, _prefix='')
A class that mimics a Weights implementation with a checkpoint file.
Unlike checkpoint-backed weights, this doesn’t carry a mapping from weight names to mmap’ed numpy arrays. Rather, when .allocate is called, this generates a backing NumPy array of the desired tensor spec on the fly and stores it. This is useful for generating weights from testing and using them in subcomponents that expect a weights implementation backed by a checkpoint.
allocate()
allocate(dtype=None, shape=None, quantization_encoding=None, device=cpu:0)
Creates a Weight that can be added to a graph.
allocated_weights
Gets the values of all weights that were allocated previously.
data()
data()
Returns data loaded from the weights at the current prefix.
-
Raises:
-
KeyError if the current prefix isn't present in the checkpoint. –
-
Return type:
exists()
exists()
Returns whether a weight with this exact name exists.
-
Return type:
items()
items()
Iterate through all allocable weights that start with the prefix.
name
property name: str
The current weight name or prefix.
raw_tensor()
raw_tensor()
Returns the numpy tensor corresponding to this weights object.
SafetensorWeights
class max.graph.weights.SafetensorWeights(filepaths, *, tensors=None, tensors_to_file_idx=None, prefix='', allocated=None, _st_weight_map=None)
Helper for loading weights into a graph.
A weight (max.graph.Weight) is tensors in a graph which are backed by an external buffer or mmap. Generally weights are used to avoid recompiling the graph when new weights are used (like from finetuning). For large-enough constants, it might be worth using weights for fast compilation times but the graph may be less optimized.
Weight classes can be used to help with graph weight allocation and naming. This protocol defines getter methods __getattr__ and __getitem__ to assist with defining names. For example, weights.a.b[1].c.allocate(…) creates a weight with the name “a.b.1.c”.
-
Parameters:
allocate()
allocate(dtype=None, shape=None, quantization_encoding=None, device=cpu:0)
Creates a Weight that can be added to a graph.
allocate_as_bytes()
allocate_as_bytes(dtype=None)
Create a Weight that can be added to the graph. Has a uint8 representation, instead of the original data type. Last dimension of the scale gets scaled by number of bytes it takes to represent the original data type. For example, [512, 256] float32 weights become [512, 1024] uint8 weights. Scalar weights will be interpreted as weights with shape [1].
allocated_weights
property allocated_weights: dict[str, ndarray[Any, dtype[_ScalarType_co]]]
Gets the values of all weights that were allocated previously.
data()
data()
Returns data loaded from the weights at the current prefix.
-
Raises:
-
KeyError if the current prefix isn't present in the checkpoint. –
-
Return type:
exists()
exists()
Returns whether a weight with this exact name exists.
-
Return type:
items()
items()
Iterate through all allocable weights that start with the prefix.
name
property name: str
The current weight name or prefix.
raw_tensor()
raw_tensor()
Returns the numpy tensor corresponding to this weights object.
WeightData
class max.graph.weights.WeightData(data, name, dtype, shape, quantization_encoding=None)
Data loaded from a checkpoint.
-
Parameters:
astype()
astype(dtype)
-
Parameters:
-
dtype (
DType) -
Return type:
data
dtype
dtype: DType
from_numpy()
classmethod from_numpy(arr, name)
name
name: str
quantization_encoding
quantization_encoding: QuantizationEncoding | None = None
shape
shape: Shape
view()
view(dtype)
-
Parameters:
-
dtype (
DType) -
Return type:
Weights
class max.graph.weights.Weights(*args, **kwargs)
Helper for loading weights into a graph.
A weight (max.graph.Weight) is tensors in a graph which are backed by an external buffer or mmap. Generally weights are used to avoid recompiling the graph when new weights are used (like from finetuning). For large-enough constants, it might be worth using weights for fast compilation times but the graph may be less optimized.
Weight classes can be used to help with graph weight allocation and naming. This protocol defines getter methods __getattr__ and __getitem__ to assist with defining names. For example, weights.a.b[1].c.allocate(…) creates a weight with the name “a.b.1.c”.
allocate()
allocate(dtype=None, shape=None, quantization_encoding=None, device=cpu:0)
Creates a Weight that can be added to a graph.
allocated_weights
property allocated_weights: dict[str, ndarray[Any, dtype[_ScalarType_co]]]
Gets the values of all weights that were allocated previously.
data()
data()
Returns data loaded from the weights at the current prefix.
-
Raises:
-
KeyError if the current prefix isn't present in the checkpoint. –
-
Return type:
exists()
exists()
Returns whether a weight with this exact name exists.
-
Return type:
items()
items()
Iterate through all allocable weights that start with the prefix.
name
property name: str
The current weight name or prefix.
raw_tensor()
raw_tensor()
Returns the numpy tensor corresponding to this weights object.
WeightsFormat
class max.graph.weights.WeightsFormat(value, names=<not given>, *values, module=None, qualname=None, type=None, start=1, boundary=None)
gguf
gguf = 'gguf'
pytorch
pytorch = 'pytorch'
safetensors
safetensors = 'safetensors'
load_weights()
max.graph.weights.load_weights(paths)
Loads weight paths into a Weights object.
-
Parameters:
-
paths (
list[Path]) – Local paths of weight files to load. -
Returns:
-
A Weights object, with all of the associated weights loaded into a single object.
-
Raises:
-
- ValueError – If an empty paths list is passed.
- ValueError – If a path provided does not exist.
-
Return type:
weights_format()
max.graph.weights.weights_format(weight_paths)
Retrieve the format of the weights files in the provided paths.
-
Parameters:
-
weight_paths (
list[Path]) – A list of file paths, containing the weights for a single model. -
Returns:
-
A WeightsFormat enum, representing whether the weights are in gguf, safetensors or pytorch format.
-
Raises:
-
ValueError – If weights type cannot be inferred from the paths.
-
Return type:
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!