mojolearn.Mamba2BlockInference

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

Bases: _RecurrentBlockInference, Mamba2Block

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.

__init__(weights, *, dt_limit=(0.0, inf))

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.