mojolearn.Mamba1Block#

class mojolearn.Mamba1Block(weights)[source]#

Bases: _MambaBase

One Mamba-1 block: RMSNorm, selective-scan mixer and residual.

Matches Hugging Face’s MambaBlock.forward. Takes (B, L, d_model) float32 and returns (B, L, d_model) float32, with the recurrent state explicit and caller-owned (Mamba1State).

Constructor keyword numeric_mode selects ‘identical’ (the default) or ‘fast’. Inputs, weights and state must be float32: bfloat16, float16 and float64 are refused by name rather than converted, and non-finite values are refused by name. Arrays are read through the buffer protocol, so NumPy is not required; arrays the block allocates are mojolearn.Array, zero-copy under numpy.asarray. There is no initializer: weights are passed in, so a cross-check can give both sides the same values. Projection weights may also be packed with mojolearn.lowbit; they are expanded exactly to float32 (see weight_format).

Fixed by the block: d_state 16, d_conv 4, expand 2, dt_rank ceil(d_model / 16), convolution bias on, projection bias off, SiLU activation.

Parameters:

weights (dict of str to float32 buffer) – Weights under the reference parameter names, with d_inner = 2 * d_model and dt_rank = ceil(d_model / 16):

  • norm.weight (d_model,)

  • in_proj.weight (2 * d_inner, d_model)

  • conv1d.weight (d_inner, 1, 4) or (d_inner, 4)

  • conv1d.bias (d_inner,)

  • x_proj.weight (dt_rank + 32, d_inner)

  • dt_proj.weight (d_inner, dt_rank)

  • dt_proj.bias (d_inner,)

  • A_log (d_inner, 16)

  • D (d_inner,)

  • out_proj.weight (d_model, d_inner)

Variables:

  • d_model (int) – Model dimension, read from norm.weight.

  • d_inner (int) – Inner dimension, 2 * d_model.

  • dt_rank (int) – Rank of the dt projection.

  • weight_format (str) – Format the weights were given in.

Examples

>>> import numpy as np, mojolearn
>>> block = mojolearn.Mamba1Block(weights)   # dict of float32 arrays, as above
>>> x = np.random.default_rng(0).random((2, 16, block.d_model), dtype=np.float32)
>>> y = block.forward(x)
>>> state = block.allocate_state(2)
>>> y = block.forward(x, state)
>>> y_next = block.step(x[:, -1], state)
__init__(weights)[source]#

Methods

allocate_state(batch_size)

Return a zero Mamba1State for batch_size sequences.

backward(x, grad_output)

Gradients of a zero-state prefill (a vector-Jacobian product).

decode_session(state)

Open a Mamba1DecodeSession for fast token-by-token decoding.

forward(x[, state, lengths])

Run the block over a sequence.

get_params([deep])

numeric_mode_used()

The numeric mode this estimator runs in.

set_params(**params)

Set constructor parameters.

step(x, state)

Run one decode token.

vendor_used()

The GPU backend this estimator runs on.
allocate_state(batch_size)[source]#

Return a zero Mamba1State for batch_size sequences.

Parameters:

batch_size (int) – Number of sequences B.

Returns:

Mamba1State

forward(x, state=None, *, lengths=None)[source]#

Run the block over a sequence.

Parameters:

  • x (float32 buffer, shape (B, L, d_model)) – Input activations: a NumPy array, array.array('f'), mojolearn.Array or any float32 buffer. Other dtypes raise TypeError.

  • state (Mamba1State or None, default None) – None runs a self-contained prefill from a zero state and discards the final state. Pass a state from allocate_state to carry it: it is read at entry and updated in place, so a later forward or step continues the sequence exactly as if the whole sequence had run in one call.

  • lengths (sequence of int or None, default None) – Makes the batch ragged: B integers in [1, L]. Row i is real at positions [0, lengths[i]) and padding after. Each real position’s output equals the row run alone at its own length, and every padding position’s output is exactly +0.0. Not allowed together with state.

Returns:

y (mojolearn.Array, shape (B, L, d_model)) – Block output, residual included. numpy.asarray(y) is zero-copy.

backward(x, grad_output)[source]#

Gradients of a zero-state prefill (a vector-Jacobian product).

Recomputes the forward with the current weights on every call; no earlier forward or cached state is used or changed. Runs under numeric_mode ‘identical’ or ‘fast’; other modes raise NotImplementedError. Gradients with respect to a carried state, decode backward and final-state cotangents are not available.

Parameters:

  • x (float32 buffer, shape (B, L, d_model)) – Forward input.

  • grad_output (float32 buffer, shape (B, L, d_model)) – Gradient of the loss with respect to the block output. Must be finite.

Returns:

grads (dict of str to mojolearn.Array) – 'x' plus one entry per constructor weight, under the same names.

step(x, state)[source]#

Run one decode token.

The same computation as forward at L = 1 with the state carried, so decoding token by token gives the same bits as prefilling the whole sequence.

Parameters:

  • x (float32 buffer, shape (B, 1, d_model) or (B, d_model)) – One token per row.

  • state (Mamba1State) – Required; updated in place.

Returns:

y (mojolearn.Array, shape (B, 1, d_model))

decode_session(state)[source]#

Open a Mamba1DecodeSession for fast token-by-token decoding.

On a GPU the session keeps the weights and the state on the device across steps; its step gives the same bytes as Mamba1Block.step. The session copies the weights and the state when it opens:

  • until close() or sync_state(), the state’s arrays are stale, and forward/step on this block refuse that state by name;

  • weight edits after opening are not seen; close and open a new session;

  • edits to the state arrays are not seen until load_state().

Parameters:

state (Mamba1State) – State to decode from.

Returns:

Mamba1DecodeSession