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).

Mojo function

mamba2_ssd_chunk_scan_varlen_fwd_shape

def mamba2_ssd_chunk_scan_varlen_fwd_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], dt_bias: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=dt_bias.static_spec], initial_states: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=initial_states.static_spec], query_start_loc: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=query_start_loc.static_spec], has_initial_state: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=has_initial_state.static_spec]) -> IndexList[Int(3)]

Computes the output shape for the mamba2_ssd_chunk_scan_varlen_fwd graph op.

Parameters:

• ​dtype (DType): Element type of the SSM input and parameter tensors (inferred).

Args:

• ​x (ManagedTensorSlice[IOSpec[_, _].Input, static_spec=x.static_spec]): Packed input tensor of shape (total_len, nheads, head_dim).
• ​dt (ManagedTensorSlice[IOSpec[_, _].Input, static_spec=dt.static_spec]): Per-head time deltas of shape (total_len, nheads).
• ​A (ManagedTensorSlice[IOSpec[_, _].Input, static_spec=A.static_spec]): Per-head scalar decay of shape (nheads,).
• ​B (ManagedTensorSlice[IOSpec[_, _].Input, static_spec=B.static_spec]): Grouped input projection of shape (total_len, ngroups, dstate).
• ​C (ManagedTensorSlice[IOSpec[_, _].Input, static_spec=C.static_spec]): Grouped output projection of shape (total_len, ngroups, dstate).
• ​D (ManagedTensorSlice[IOSpec[_, _].Input, static_spec=D.static_spec]): Per-head skip connection of shape (nheads,); may be empty when unused.
• ​dt_bias (ManagedTensorSlice[IOSpec[_, _].Input, static_spec=dt_bias.static_spec]): Per-head bias added to dt of shape (nheads,); may be empty when unused.
• ​initial_states (ManagedTensorSlice[IOSpec[_, _].Input, static_spec=initial_states.static_spec]): Optional initial SSM states of shape (batch, nheads, head_dim, dstate) in float32; may be empty when has_initial_state is all false.
• ​query_start_loc (ManagedTensorSlice[IOSpec[_, _].Input, static_spec=query_start_loc.static_spec]): Cumulative sequence lengths of shape (batch + 1,) in int32.
• ​has_initial_state (ManagedTensorSlice[IOSpec[_, _].Input, static_spec=has_initial_state.static_spec]): Per-sequence flag of shape (batch,) in bool indicating whether to load initial_states; may be empty when no initial states are used.

Returns:

IndexList[Int(3)]