Python class
WeightData
WeightData
class max.graph.weights.WeightData(data, name, dtype, shape, quantization_encoding=None)
Bases: DLPackArray
Container for weight tensor data with metadata.
WeightData encapsulates a weight tensor along with its metadata,
providing utilities for type conversion and format compatibility.
It supports the DLPack protocol for efficient tensor sharing between
frameworks.
-
Parameters:
-
- data (DLPackArray)
- name (str)
- dtype (DType)
- shape (Shape)
- quantization_encoding (QuantizationEncoding | None)
astype()
astype(dtype)
Convert the weight data to a different dtype.
This method performs actual data conversion of the underlying tensor data. Special handling is provided for bfloat16 conversions using PyTorch when available.
During cross-compilation (warm-cache) scenarios where no GPU is available, this method skips the actual data conversion but still returns a WeightData with the target dtype. This is safe because the weight data won’t be used for inference during compilation - only the dtype metadata matters for graph construction.
weight_data = weights.model.layer.weight.data()
fp16_data = weight_data.astype(DType.float16)-
Parameters:
-
dtype (DType) – Target data type for conversion.
-
Returns:
-
A new WeightData instance with the converted data.
-
Return type:
data
data: DLPackArray
The weight tensor as a DLPack array.
dtype
dtype: DType
Data type of the tensor (for example, DType.float32, DType.uint8).
from_numpy()
classmethod from_numpy(arr, name)
Create WeightData from a numpy array.
name
name: str
Hierarchical name of the weight (for example, model.layers.0.weight).
quantization_encoding
quantization_encoding: QuantizationEncoding | None = None
Optional quantization scheme applied to the weight.
shape
shape: Shape
Shape of the tensor as a Shape object.
to_buffer()
to_buffer()
Mutates the data into a Buffer.
-
Return type:
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!