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

gated_delta_recurrence_fwd_shape

def gated_delta_recurrence_fwd_shape[work_dtype: DType, state_dtype: DType](qkv_conv_output: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=qkv_conv_output.static_spec], decay_per_token: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=decay_per_token.static_spec], beta_per_token: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=beta_per_token.static_spec], recurrent_state: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=recurrent_state.static_spec], slot_idx: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=slot_idx.static_spec], input_row_offsets: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=input_row_offsets.static_spec]) -> IndexList[Int(2)]

Computes the output shape for the gated_delta_recurrence_fwd graph op.

Parameters:

• ​work_dtype (DType): Element type of the per-token input tensors (inferred).
• ​state_dtype (DType): Element type of the persistent recurrent-state pool, independent of work_dtype (inferred).

Args:

• ​qkv_conv_output (ManagedTensorSlice[IOSpec[_, _].Input, static_spec=qkv_conv_output.static_spec]): Ragged conv output of shape [total_seq_len, conv_dim].
• ​decay_per_token (ManagedTensorSlice[IOSpec[_, _].Input, static_spec=decay_per_token.static_spec]): Per-token decay factors of shape [total_seq_len, num_value_heads].
• ​beta_per_token (ManagedTensorSlice[IOSpec[_, _].Input, static_spec=beta_per_token.static_spec]): Per-token beta gates of shape [total_seq_len, num_value_heads].
• ​recurrent_state (ManagedTensorSlice[IOSpec[_, _].Input, static_spec=recurrent_state.static_spec]): Mutable recurrent-state pool of shape [max_slots, num_value_heads, key_head_dim, value_head_dim].
• ​slot_idx (ManagedTensorSlice[IOSpec[_, _].Input, static_spec=slot_idx.static_spec]): Per-batch slot indices into the recurrent-state pool, shape [batch_size].
• ​input_row_offsets (ManagedTensorSlice[IOSpec[_, _].Input, static_spec=input_row_offsets.static_spec]): Cumulative row offsets per batch, shape [batch_size + 1].

Returns:

IndexList[Int(2)]