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

forward_sequential_layers

forward_sequential_layers()

max.nn.transformer.forward_sequential_layers(layers, *, inputs_for_layer, initial_hidden_states, on_layer_output=None, subgraph_layer_groups=None, name_for_subgraph=<function <lambda>>, weight_prefix_for_layer=None, layer_ids=None)

source

Forward pass through sequential layers with optional subgraph groups.

Parameters:

• layers (Sequence[Module]) – The sequence of layers to forward through.

• inputs_for_layer (Callable[[int, list[TensorValue]], Sequence[Tree[Any]]]) – A callable that takes the layer index and previous hidden states and returns a list of values.

• initial_hidden_states (list[TensorValue]) – The hidden states to provide to the first layer.

• on_layer_output (Callable[[int, list[TensorValue]], None] | None) – An optional callable invoked after each layer with the layer index and output hidden states.

• subgraph_layer_groups (list[list[int]] | None) –

  A sequence of layer index groups. Each group is a sequence of layer indices that should share a single compiled subgraph. If None or empty, all layers are called directly without subgraphs. Layers not listed in any group also fall through to direct call.

  Note

  All layers in a group must share the same operand types. The subgraph signature is fixed at the first layer in the group; other layers calling it with different shapes fail with an MLIR Subgraph ... has wrong type for argument error. Architectures with per-layer-variable shapes (variable head count, mixed dense and sparse MLP) must group layers by signature.

• name_for_subgraph (Callable[[int], str]) – A callable that takes the group index and returns the subgraph name.

• weight_prefix_for_layer (Callable[[int], str] | None) – A callable that takes the layer index and returns the weight prefix string for that layer. Required when using subgraphs.

• layer_ids (Sequence[int] | None) – The index of each of layers that the callbacks receive and subgraph_layer_groups lists, such as its index in the whole model when layers are one pipeline stage’s. Defaults to each layer’s position in layers.

Returns:

The output hidden states of the last layer.

Return type:

list[TensorValue]