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_scan_fwd_shape

def varlen_selective_scan_fwd_shape[dtype: DType](u: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=u.static_spec], delta: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=delta.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], delta_bias: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=delta_bias.static_spec], query_start_loc: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=query_start_loc.static_spec], cache_indices: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=cache_indices.static_spec], has_initial_state: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=has_initial_state.static_spec]) -> IndexList[Int(2)]

Returns the output shape for the varlen_selective_scan_fwd op.

The variable-length selective scan output has the same shape as the packed input u: (dim, total_length).

Parameters:

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

Args:

u (ManagedTensorSlice[IOSpec[_, _].Input, static_spec=u.static_spec]): Packed input tensor with shape (dim, total_length).
delta (ManagedTensorSlice[IOSpec[_, _].Input, static_spec=delta.static_spec]): Time-step tensor with shape (dim, total_length).
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 projection with shape (ngroups, dstate, total_length).
C (ManagedTensorSlice[IOSpec[_, _].Input, static_spec=C.static_spec]): Output projection with shape (ngroups, dstate, total_length).
D (ManagedTensorSlice[IOSpec[_, _].Input, static_spec=D.static_spec]): Skip connection with shape (dim,).
delta_bias (ManagedTensorSlice[IOSpec[_, _].Input, static_spec=delta_bias.static_spec]): Delta bias with shape (dim,).
query_start_loc (ManagedTensorSlice[IOSpec[_, _].Input, static_spec=query_start_loc.static_spec]): Cumulative sequence lengths with shape (batch + 1,).
cache_indices (ManagedTensorSlice[IOSpec[_, _].Input, static_spec=cache_indices.static_spec]): Indices into SSM states with shape (batch,).
has_initial_state (ManagedTensorSlice[IOSpec[_, _].Input, static_spec=has_initial_state.static_spec]): Boolean mask with shape (batch,).

Returns:

IndexList[Int(2)]: The output tensor shape (dim, total_length).