mojolearn.TransformerBlockInference#

class mojolearn.TransformerBlockInference(weights, *, n_heads, n_kv_heads=None, head_dim=None, window=0, rope_theta=10000.0, rope_scaling=None, rope_dim=None, max_positions=8192, qkv_bias=False, o_bias=False, norm='rmsnorm', norm_eps=1e-06, norm_bias=False, mlp='swiglu', mlp_bias=False, qk_norm=False, attn_softcap=None, numeric_profile=None)[source]#

Bases: TransformerBlock

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

Takes the same constructor arguments as TransformerBlock (the nine named weights, n_heads, n_kv_heads, head_dim, window) and applies the same weight checks. Supports forward(x, lengths=None) from a zero state, allocate_state(B, max_tokens), 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. A block has no output head, so it returns hidden states, not logits. Runs on the host binding shipped with mojolearn, loaded on first use.

Examples

>>> blk = mojolearn.TransformerBlockInference(weights, n_heads=4)
>>> y = blk.forward(x)

__init__(weights, *, n_heads, n_kv_heads=None, head_dim=None, window=0, rope_theta=10000.0, rope_scaling=None, rope_dim=None, max_positions=8192, qkv_bias=False, o_bias=False, norm='rmsnorm', norm_eps=1e-06, norm_bias=False, mlp='swiglu', mlp_bias=False, qk_norm=False, attn_softcap=None, numeric_profile=None)#

Methods

`allocate_state`(batch_size, max_tokens)	A zero `TransformerState` for `batch_size` sequences of up to `max_tokens` positions each.
`backward`(x, grad_output)	Not available: inference is forward only.
`decode_session`(state)	Open a `TransformerDecodeSession` on this block and `state`.
`forward`(x[, state, lengths])	Run the block on `(B, L, d_model)` float32 input and return the output, same shape.
`get_params`([deep])
`numeric_mode_used`()	The numeric mode this estimator runs in.
`set_params`(**params)	Set constructor parameters.
`step`(x, state)	Decode one token per sequence, continuing the cache in `state`.
`vendor_used`()	The GPU backend this estimator runs on.

backward(x, grad_output)[source]#: Not available: inference is forward only. Always raises NotImplementedError; train with the GPU block class.