mojolearn.Mamba2State#

class mojolearn.Mamba2State(conv_window, h, buffer_xbc, buffer_dtraw, buffered_tokens=0)[source]#

Bases: object

Recurrent state of a Mamba2Block, owned by the caller.

With H = d_model / 32 heads and CD = 2 * d_model + 256 convolution channels. Make one with Mamba2Block.allocate_state; it is all zeros (and buffered_tokens == 0) before the first token. forward and step update it in place. Setting h nonzero on a fresh state supplies the reference’s initial_states.

Variables:

  • conv_window (float32 buffer, shape (B, CD, 4)) – The last four pre-convolution xBC inputs, oldest first.

  • h (float32 buffer, shape (B, H, 64, 128)) – SSM state as of the last completed 256-token chunk.

  • buffer_xbc (float32 buffer, shape (B, 256, CD)) – Post-convolution, post-SiLU xBC rows of the open chunk.

  • buffer_dtraw (float32 buffer, shape (B, 256, H)) – Raw dt rows of the open chunk.

  • buffered_tokens (int) – Number of valid rows in the open-chunk buffers, in [0, 256). Decoding resumes the open chunk from these rows.

__init__(conv_window, h, buffer_xbc, buffer_dtraw, buffered_tokens=0)[source]#