Skip to content

base

base

Inference job for running inference on trained models.

Provides an inference job that either owns its state or borrows it from a trainer.

InferenceJob(spec: ExecutionSpec, base: Optional[Node] = None)

Bases: CheckpointedJob[C], LoggingJob[C], Generic[C, M]

Run inference with owned state or with state borrowed from a trainer.

A normally constructed job owns its node, state, store, and clock. Jobs made with :meth:from_trainer share the trainer's store and dynamically read its current node and state. The trainer remains the sole owner of that clock, so calling :meth:tick on the borrowed job raises RuntimeError.

Example subclass

class GPTInference(InferenceJob): MODEL = GPT

@staticmethod
def forward(state, params, batch, key, deterministic):
    # Custom forward implementation
    ...

Attributes:

Name Type Description
state TrainState

TrainState with params
mesh Mesh

JAX device mesh
state_sharding PyTree[NamedSharding]

NamedSharding
replicas int

Global replica count.
local_replicas int

Replica count on the current host.
per_device_batch_size int

Local examples per replica.
block_size int

Maximum context length.
key

PRNG key

state: train_state.TrainState property writable

Return the owned state or the trainer's current state.

template: PyTree[Any] property

Return a params-only, sharded shape tree for partial restoration.

from_trainer(trainer: BaseTrainer[Any, Any]) -> Self classmethod

Create an inference view over a trainer's current state and node.

tick() -> None

Advance an owned node; trainer-owned inference uses the trainer's clock.

initialize() -> None

Initialize inference state from the configured model.

surgery(partial: PyTree[bool]) -> PyTree[Any]

Initialize only parameter leaves missing from a restored checkpoint.

apply(state: PyTree[Any], metadata: ValueRow) -> None

Build an inference TrainState from restored model parameters.

forward(state: train_state.TrainState, params: Any, batch: Tuple[jax.Array, Optional[jax.Array], jax.Array], key: Optional[jax.Array] = None, deterministic: bool = False, mutable: Optional[list[str] | tuple[str, ...]] = None, extra_variables: Optional[dict[str, Any]] = None, cache_max_len: Optional[int] = None) -> Any staticmethod

Forward pass with optional mutable variable collections (e.g. KV cache).

Parameters:

Name Type Description Default
mutable Optional[list[str] | tuple[str, ...]]

List of mutable variable collections (e.g. ['cache']). When provided, returns ((logits, loss), mutated_variables).

 None
extra_variables Optional[dict[str, Any]]

Additional variable collections to pass alongside params (e.g. {'cache': cache_state} for decode steps).

 None
cache_max_len Optional[int]

Forwarded to model __call__ so attention layers size their KV cache to actual decode need rather than the model's full block_size. Only forwarded when not None, so models that don't accept it are unaffected.

 None

Returns:

Type Description
Any

(logits, loss, meta) when mutable is None.
Any

((logits, loss, meta), mutated_variables) when mutable is provided.

pad(seqs: List[List[int]], pad_token: int = 0, pad_to: Optional[int] = None) -> Tuple[np.ndarray, np.ndarray] staticmethod

Left-pad sequences to uniform length.

Parameters:

Name Type Description Default
seqs List[List[int]]

List of token id lists

 required
pad_token int

Token to use for padding (default 0)

 0
pad_to Optional[int]

Minimum length to pad to (default None, uses max seq length)

 None

Returns:

Name Type Description
padded ndarray

(batch_size, max_len) host int32 array
mask ndarray

(batch_size, max_len) host bool array, True for real tokens

rollout(inputs: List[Union[str, ChatTemplate, jax.Array, List[int]]], encoding: Optional[Tokenizer] = None, max_new_tokens: Optional[int] = None, max_prompt_length: Optional[int] = None, temperature: float = 0.0, top_p: float = 1.0, chunk_size: int = 200, return_type: Literal['decoded', 'indices', 'output_decoded', 'output_indices', 'raw_indices'] = 'decoded') -> Union[List[Union[str, ChatTemplate]], List[str], List[List[int]]]

Autoregressive rollout of the language model.

Parameters:

Name Type Description Default
inputs List[Union[str, ChatTemplate, Array, List[int]]]

List of raw strings, ChatTemplates, or pre-tokenized 1D jax arrays of token ids.

 required
encoding Optional[Tokenizer]

Tokenizer for encoding/decoding. Required when any input is a string/ChatTemplate or when return_type is one of the decoded variants. May be omitted when all inputs are pre-tokenized arrays AND return_type is an indices variant.

 None
max_new_tokens Optional[int]

Maximum number of new tokens to generate. Defaults to block_size - max_prompt_length (or block_size // 2 if max_prompt_length is also unset).

 None
max_prompt_length Optional[int]

Length to which prompts are padded. Defaults to block_size - max_new_tokens. Inputs longer than this raise. Holding this fixed across calls keeps the JIT trace stable — varying it triggers recompiles.

 None
temperature float

Sampling temperature (0.0 for greedy).

 0.0
top_p float

Nucleus sampling threshold.

 1.0
chunk_size int

Number of batches per JIT chunk.

 200
return_type
  • "decoded": full prompt + generated tokens, decoded, with left-pad stripped.
  • "indices": full prompt + generated tokens as ids, with left-pad stripped.
  • "output_decoded": generated portion only, decoded.
  • "output_indices": generated portion only as ids.
  • "raw_indices": full fixed-shape rows, left padding preserved. Shape is logically (N, max_prompt_length + max_new_tokens).