Study Plan

Mamba 2 Gated State Scan

Language ModelsLinear AlgebraNeural Networks

Medium

Implement a gated recurrent state scan. At position $t$ , $\gamma_t$ is a scalar decay gate, $q_t$ , $k_t$ , and $v_t$ are query, key, and value vectors, $S_t$ is the state, and $y_t$ is the output. Superscript T denotes transpose.

S_t=\gamma_tS_{t-1}+k_tv_t^{\mathsf T}.

y_t=q_t^{\mathsf T}S_t.

Each batch state starts at zero, and the current output reads the updated state. Return a dictionary containing outputs with shape (B, S, Dv) and final_state with shape (B, Dk, Dv), both PyTorch tensors with the same dtype and device as q.

Loading visualization...

Examples

Input: q = [[[1,0],[0,1]]], k = [[[1,0],[0,1]]], v = [[[2,3],[4,5]]], gamma = [[0.5,1]]

Output: {"outputs":[[[2,3],[4,5]]],"final_state":[[[2,3],[4,5]]]}

Explanation: The first key writes the first state row. The second gate preserves that row while the orthogonal key writes the second row.

Input: q = [[[1,1],[2,0]]], k = [[[1,0],[0,1]]], v = [[[3],[5]]], gamma = [[0,0]]

Output: {"outputs":[[[3],[0]]],"final_state":[[[0],[5]]]}

Input: q = [[[1]]], k = [[[2]]], v = [[[3]]], gamma = [[0.25]]

Output: {"outputs":[[[6]]],"final_state":[[[6]]]}

Hint 1

torch.unsqueeze adds singleton dimensions for broadcasting a scalar gate.

Hint 2

torch.bmm handles the per-batch matrix products.

Requirements

Use PyTorch, with float32 intermediate calculations and results cast to the query dtype.
Each batch item has its own independent recurrent state.
Do not modify any input tensor.

Constraints

q and k have shape (B, S, Dk), v has shape (B, S, Dv), and gamma has shape (B, S).
B is batch size, S is sequence length, and Dk and Dv are key and value widths; all are positive.
All inputs are PyTorch tensors on the same device; q, k, and v share dtype float32 or bfloat16, while gamma is float32.
Every gamma value lies in [0, 1].
Inputs contain finite values and may be noncontiguous.

Try Similar Problems

Parallel and Recurrent Linear AttentionMedium Gated DeltaNet State UpdateHard Top-k MoE Router with Load StatisticsHard Orthogonalize a Muon Momentum UpdateEasy RMSNorm Forward PassEasy Train a BPE Tokenizer From Scratch in PythonMedium Canonicalize Pairwise Preference DataMedium Compute Temperature-Based Data MixturesMedium Scaled Dot-Product Attention in PythonMedium Jacobian of a Neural Network LayerHard JVP and VJP for a 2-Layer NetworkHard Hessian Tensor of a Vector-Valued FunctionMedium

Mamba 2 Gated State Scan

Language ModelsLinear AlgebraNeural Networks

Medium

S_t=\gamma_tS_{t-1}+k_tv_t^{\mathsf T}.

y_t=q_t^{\mathsf T}S_t.

Loading visualization...

Examples

Input: q = [[[1,0],[0,1]]], k = [[[1,0],[0,1]]], v = [[[2,3],[4,5]]], gamma = [[0.5,1]]

Output: {"outputs":[[[2,3],[4,5]]],"final_state":[[[2,3],[4,5]]]}

Explanation: The first key writes the first state row. The second gate preserves that row while the orthogonal key writes the second row.

Input: q = [[[1,1],[2,0]]], k = [[[1,0],[0,1]]], v = [[[3],[5]]], gamma = [[0,0]]

Output: {"outputs":[[[3],[0]]],"final_state":[[[0],[5]]]}

Input: q = [[[1]]], k = [[[2]]], v = [[[3]]], gamma = [[0.25]]

Output: {"outputs":[[[6]]],"final_state":[[[6]]]}

Hint 1

torch.unsqueeze adds singleton dimensions for broadcasting a scalar gate.

Hint 2

torch.bmm handles the per-batch matrix products.

Requirements

Use PyTorch, with float32 intermediate calculations and results cast to the query dtype.
Each batch item has its own independent recurrent state.
Do not modify any input tensor.

Constraints

q and k have shape (B, S, Dk), v has shape (B, S, Dv), and gamma has shape (B, S).
B is batch size, S is sequence length, and Dk and Dv are key and value widths; all are positive.
All inputs are PyTorch tensors on the same device; q, k, and v share dtype float32 or bfloat16, while gamma is float32.
Every gamma value lies in [0, 1].
Inputs contain finite values and may be noncontiguous.

Try Similar Problems

q =

Accepts: array

k =

Accepts: array

v =

Accepts: array

gamma =

Accepts: array