mojolearn.Mamba1BlockInference#
- class mojolearn.Mamba1BlockInference(weights)[source]#
Bases:
_RecurrentBlockInference,Mamba1BlockForward-only
Mamba1Blockon the CPU, from weights trained on a GPU.Takes the same constructor arguments as
Mamba1Block(its ten named weights) and applies the same weight checks. Supportsforward(x, lengths=None)from a zero state,allocate_state(B),forward(x, state)andstep(x, state); decoding one token at a time with a carried state gives the same bits as one forward pass over the whole sequence.backwardraisesNotImplementedError. 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)#
Methods
allocate_state(batch_size)Return a zero
Mamba1Stateforbatch_sizesequences.backward(x, grad_output)Gradients of a zero-state prefill (a vector-Jacobian product).
decode_session(state)Open a
Mamba1DecodeSessionfor 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.