mojolearn.Mamba3State#

class mojolearn.Mamba3State(theta, h, buffer_qrot, buffer_krot, buffer_v, buffer_dt, buffer_sig, buffer_adt, pending_k, pending_v, buffered_tokens=0, pending=False)[source]#

Bases: object

Recurrent state of a Mamba3Block, owned by the caller.

With H = d_model / 32 heads, N = 128 (QK head dimension), P = 64 (V head dimension), R = 32 rotary angles and Q = 64 (the Mamba-3 chunk size). Make one with Mamba3Block.allocate_state; it is all zeros (buffered_tokens == 0, pending == False) before the first token. forward and step update it in place.

Variables:

  • theta (float32 buffer, shape (B, H, R)) – Rotary angles, each in [0, 2*pi).

  • h (float32 buffer, shape (B, H, P, N)) – SSM state entering the last working chunk.

  • buffer_krot (buffer_qrot,) – Rotated, unscaled q and k rows of the working chunk.

  • buffer_v (float32 buffer, shape (B, Q, H, P)) – Raw v rows.

  • buffer_adt (buffer_dt, buffer_sig,) – dt, sigmoid(trap) and A*dt rows.

  • pending_k (float32 buffer, shape (B, H, N)) – Input-state k, used only when pending.

  • pending_v (float32 buffer, shape (B, H, P)) – Input-state v, used only when pending.

  • buffered_tokens (int) – Valid buffer rows, in [0, Q]. After the first call it is always in [1, Q].

  • pending (bool) – True when set_input_states supplied the state; consumed by the next call.

__init__(theta, h, buffer_qrot, buffer_krot, buffer_v, buffer_dt, buffer_sig, buffer_adt, pending_k, pending_v, buffered_tokens=0, pending=False)[source]#

Methods

set_input_states(theta, h, k, v)

Start from a reference Input_States continuation.
set_input_states(theta, h, k, v)[source]#

Start from a reference Input_States continuation.

Copies the given values into theta, h, pending_k and pending_v and marks them pending for the next call. The state must be fresh (no tokens yet); that is checked at the next call. Arrays must already be float32 with the shapes above; nothing is converted.

Parameters:

  • theta (float32 buffer, shape (B, H, R))

  • h (float32 buffer, shape (B, H, P, N))

  • k (float32 buffer, shape (B, H, N))

  • v (float32 buffer, shape (B, H, P))