Skip to main content

/

Docs

Nightly

v25.1

Log in

Python module

layer

Layer​

class max.pipelines.nn.layer.Layer

Base Layer class.

Currently, only functionality is for adding hooks to the call function of each layer to support testing, debugging or profiling.

LayerList​

class max.pipelines.nn.layer.LayerList(layers: Sequence[Layer])

Stores a list of layers.

Can be used as a regular python list.

append()​

append(layer: Layer)

extend()​

extend(layer: Layer)

insert()​

insert(i, layer: Layer)

sublayers​

property sublayers*: dict[str, max.pipelines.nn.layer.layer.LayerV2]*

LayerV2​

class max.pipelines.nn.layer.LayerV2

(new) Base class for model layers with weight and device management.

This will be merged with the above class once all layers have been moved to V2.

layer_weights​

property layer_weights*: dict[str, max.graph.weight.Weight]*

load_state_dict()​

load_state_dict(state_dict: Mapping[str, DLPackArray | ndarray | WeightData], *, override_quantization_encoding: bool = False) → None

Sets the values of all weights in this model.

  • Parameters:

    • state_dict – A map from weight name to a numpy array or max.driver.Tensor.
    • override_quantization_encoding – Whether to override the weight quantization based on the loaded value.

  • Raises:

    Error if any weight in the model is not present in the state dict. –

raw_state_dict()​

raw_state_dict() → dict[str, max.graph.weight.Weight]

Returns all weights objects in the model. Unlike state_dict, this returns max.graph.Weight objects instead of the assigned values. Some parameters inside the Weight can be configured before a graph is built. Do not change these attributes after building a graph:

  • Weight.shape
  • Weight.dtype
  • Weight.quantization_encoding
  • Weight.device
  • Weight.align

  • Returns:

    Map from weight name to the max.graph.Weight object.

set_shared_weight()​

set_shared_weight(name: str, weight: Weight)

state_dict()​

state_dict(auto_initialize: bool = True) → dict[str, Union[max.driver.tensor.DLPackArray, numpy.ndarray]]

Returns values of all weights in the model.

The values returned are the same as the values set in load_state_dict. If load_state_dict has not been called and none of the weights have values, then they are initialized to zero.

  • Parameters:

    auto_initialize – Determines whether to initialize weights to zero if the weight value has not been loaded. If this is False, a ValueError is raised if an uninitialized weight is found.

  • Returns:

    Map from weight name to the weight value (can be numpy array or max.driver.Tensor).

sublayers​

property sublayers*: dict[str, max.pipelines.nn.layer.layer.LayerV2]*

to()​

to(*devices: DeviceRef, sharding_strategy: ShardingStrategy | None = None) → None

add_layer_hook()​

max.pipelines.nn.layer.add_layer_hook(fn: Callable[[Layer, Tuple[Any, ...], Dict[str, Any], Any], Any]) → None

Adds a hook to call a function after each layer’s __call__.

The function will be passed four inputs: the layer, input_args, input_kwargs and outputs. The function can either return None or new outputs that will replace the layer returned outputs.

Note that input and outputs contain graph Values, which show limited information (like shape and dtype). You can still see the computed values if you include the Value in the graph.output op, or call value.print.

Example of printing debug inputs:

def print_info(layer, args, kwargs, outputs):
    print("Layer:", type(layer).__name__)
    print("Input args:", args)
    print("Input kwargs:", kwargs)
    print("Outputs:", outputs)
    return outputs

add_layer_hook(print_info)
def print_info(layer, args, kwargs, outputs):
    print("Layer:", type(layer).__name__)
    print("Input args:", args)
    print("Input kwargs:", kwargs)
    print("Outputs:", outputs)
    return outputs

add_layer_hook(print_info)

clear_hooks()​

max.pipelines.nn.layer.clear_hooks()

Remove all hooks.

recursive_named_layers()​

max.pipelines.nn.layer.recursive_named_layers(parent: LayerV2, prefix: str = '') → Iterable[tuple[str, max.pipelines.nn.layer.layer.LayerV2]]

Recursively walks through the layers and generates names.