For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /max/get-started.md).

Mojo function

varlen_selective_state_update_shape

def varlen_selective_state_update_shape[dtype: DType](x: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=x.static_spec], dt: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=dt.static_spec], A: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=A.static_spec], B: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=B.static_spec], C: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=C.static_spec], D: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=D.static_spec], z: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=z.static_spec], dt_bias: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=dt_bias.static_spec], state_batch_indices: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=state_batch_indices.static_spec]) -> Tuple[IndexList[Int(4)], IndexList[Int(3)]]

Returns the output shapes for the varlen_selective_state_update op.

The update produces the updated SSM state and the single-step output.

Parameters:

dtype (DType): Element type of the state update input and output tensors.

Args:

x (ManagedTensorSlice[IOSpec[_, _].Input, static_spec=x.static_spec]): Input tensor with shape (batch, nheads, dim).
dt (ManagedTensorSlice[IOSpec[_, _].Input, static_spec=dt.static_spec]): Time-delta tensor with shape (batch, nheads, dim).
A (ManagedTensorSlice[IOSpec[_, _].Input, static_spec=A.static_spec]): State transition matrix with shape (nheads, dim, dstate).
B (ManagedTensorSlice[IOSpec[_, _].Input, static_spec=B.static_spec]): Input matrix with shape (batch, ngroups, dstate).
C (ManagedTensorSlice[IOSpec[_, _].Input, static_spec=C.static_spec]): Output matrix with shape (batch, ngroups, dstate).
D (ManagedTensorSlice[IOSpec[_, _].Input, static_spec=D.static_spec]): Skip connection with shape (nheads, dim).
z (ManagedTensorSlice[IOSpec[_, _].Input, static_spec=z.static_spec]): Gating tensor with shape (batch, nheads, dim).
dt_bias (ManagedTensorSlice[IOSpec[_, _].Input, static_spec=dt_bias.static_spec]): Time-delta bias with shape (nheads, dim).
state_batch_indices (ManagedTensorSlice[IOSpec[_, _].Input, static_spec=state_batch_indices.static_spec]): Batch indices into the state buffer with shape (batch,).

Returns:

Tuple[IndexList[Int(4)], IndexList[Int(3)]]: A tuple (state_shape, output_shape) where state_shape is (batch, nheads, dim, dstate) and output_shape matches x.shape().