mojolearn.Mamba1State#

class mojolearn.Mamba1State(conv_window, h, *, profile_id=1, absolute_position=0, affine_boundary=None, affine_a=None, affine_b=None)[source]#

Bases: object

Recurrent state of a Mamba1Block, owned by the caller.

Make one with Mamba1Block.allocate_state; it is all zeros before the first token. forward and step update the arrays in place. The arrays are mojolearn.Array when allocate_state made them, or any writable float32 buffer (a NumPy array included) you supply; their bytes can be saved and restored exactly.

Variables:

  • conv_window (float32 buffer, shape (B, d_inner, 4)) – The last four inputs to the convolution, oldest first.

  • h (float32 buffer, shape (B, d_inner, 16)) – SSM state.

__init__(conv_window, h, *, profile_id=1, absolute_position=0, affine_boundary=None, affine_a=None, affine_b=None)[source]#