mojolearn.Mamba2BlockInference#

class mojolearn.Mamba2BlockInference(weights, *, dt_limit=(0.0, inf))#

Forward-only Mamba2Block on the CPU, from weights trained on a GPU.

Takes the same constructor arguments as Mamba2Block (its nine named weights and dt_limit, default (0.0, inf)) and applies the same weight checks. Supports forward(x, lengths=None) from a zero state, allocate_state(B), forward(x, state) and step(x, state); decoding one token at a time with a carried state gives the same bits as one forward pass over the whole sequence. backward raises NotImplementedError. Runs on the host binding shipped with mojolearn; it is loaded on first use, so an install without it still imports and raises when the class is used.

See also

User guide: Neural networks

Methods

allocate_state(batch_size)

Return a zero Mamba2State for batch_size sequences.

backward(x, grad_output)

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

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)#

Return a zero Mamba2State for batch_size sequences.

Set its h nonzero before the first call to supply initial_states.

Parameters:

batch_size (int) – Number of sequences B.

Returns:

Mamba2State

backward(x, grad_output)#

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’ (Apple GPUs only); 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.

forward(x, state=None, *, lengths=None)#

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 (Mamba2State 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.

numeric_mode_used()#

The numeric mode this estimator runs in.

Returns:

str – 'identical' or 'fast', read from the compiled binary the estimator uses, not from the argument it was given.

set_params(**params)#

Set constructor parameters. The estimator is rebuilt through its constructor, so every constructor check runs again and any fitted state is dropped, as a new setting invalidates it.

step(x, state)#

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 (Mamba2State) – Required; updated in place.

Returns:

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

vendor_used()#

The GPU backend this estimator runs on.

Returns:

str – 'metal', 'cuda' or 'hip', read from the compiled binary the estimator uses.