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

selective_scan_update_shape

def selective_scan_update_shape[dtype: DType](state_in: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=state_in.static_spec], 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]) -> Tuple[IndexList[Int(3)], IndexList[Int(2)]]

Returns the output shapes for the selective_scan_update op.

The update step produces two tensors: the updated SSM state and the single-step output for the new token.

Args:

state_in (ManagedTensorSlice[IOSpec[_, _].Input, static_spec=state_in.static_spec]): Input SSM state with shape (batch, dim, dstate).
x (ManagedTensorSlice[IOSpec[_, _].Input, static_spec=x.static_spec]): Input token with shape (batch, dim).
dt (ManagedTensorSlice[IOSpec[_, _].Input, static_spec=dt.static_spec]): Time-delta tensor with shape (batch, dim).
A (ManagedTensorSlice[IOSpec[_, _].Input, static_spec=A.static_spec]): State transition matrix with shape (dim, dstate).
B (ManagedTensorSlice[IOSpec[_, _].Input, static_spec=B.static_spec]): Input matrix with shape (batch, n_groups, dstate).
C (ManagedTensorSlice[IOSpec[_, _].Input, static_spec=C.static_spec]): Output matrix with shape (batch, n_groups, dstate).
D (ManagedTensorSlice[IOSpec[_, _].Input, static_spec=D.static_spec]): Skip connection with shape (dim,).
z (ManagedTensorSlice[IOSpec[_, _].Input, static_spec=z.static_spec]): Gating tensor with shape (batch, dim).
dt_bias (ManagedTensorSlice[IOSpec[_, _].Input, static_spec=dt_bias.static_spec]): Time-delta bias with shape (dim,).

Returns:

Tuple[IndexList[Int(3)], IndexList[Int(2)]]: A tuple of (state_out_shape, output_shape) where state_out_shape matches state_in.shape() and output_shape matches x.shape().