Skip to content

Index

experiments

Evaluator(spec: ExecutionSpec, base: Node | None = None)

Bases: InferenceJob[EvaluatorConfig, M], Generic[M]

InferenceJob that runs evaluations and saves results.

Created from a trainer, holds a list of evaluations, and runs them when run() is called. Standalone inference subclasses can use the evaluation components directly.

Example

evaluator = Evaluator.from_trainer(trainer) evaluator() # Runs evaluations and saves results

config(evaluations: Optional[List[type[Evaluation]]] = None) -> List[type[Any]] classmethod

Return evaluator and child-evaluation configuration schemas.

Parameters:

Name Type Description Default
evaluations Optional[List[type[Evaluation]]]

Evaluation classes to aggregate. Defaults to the evaluator's static EVALUATION declaration.

 None

Returns:

Type Description
List[type[Any]]

Schemas required to construct the evaluator and its children.

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

Create Evaluator from trainer.

Parameters:

Name Type Description Default
trainer BaseTrainer[Any, Any]

BaseTrainer instance to get inference state from

 required

Returns:

Type Description
Evaluator[M]

Evaluator instance ready to run evaluations

evaluate(reduce: str = 'mean', return_intermediates: bool = False, **kwargs: Any) -> Any

Run all evaluations.

Parameters:

Name Type Description Default
reduce str

passed through to each evaluation. "mean"/"sum" → float per evaluation; "none" → np.ndarray of per-sample scores.

 'mean'
return_intermediates bool

when True, also return the per-evaluation list of (x, mask) rollouts (one inner list per evaluation).

 False
**kwargs Any

forwarded to each evaluation's call (e.g. temperature, top_p, chunk_size).

 {}

run() -> None

Run all evaluations and save results to disk.

Evaluation

Bases: ABC

Abstract base class for all evaluations.

name: str abstractmethod property

Name of this evaluation.

config() -> List[type[Any]] classmethod

Return configuration schemas required by this evaluation.

Returns:

Type Description
List[type[Any]]

The evaluation's schema, or an empty list when it has none.

prefix() -> str

Prefix for metrics from this evaluation.

__len__() -> int abstractmethod

Number of samples in this evaluation.

__call__(inference: InferenceJob[Any, M], encoding: Any, reduce: str = 'mean', return_intermediates: bool = False, **kwargs: Any) -> Any abstractmethod

Run the evaluation and return a score (and optionally intermediates).

When return_intermediates=True, also returns a list of (x, padding_mask) numpy arrays — one per sample — available on every host so that an RL trainer can use them as a training batch.

score(*args: Any) -> List[float]

Return one float per evaluation sample. Subclasses override.

find_accumulation_steps(dataset_size: int, max_batch_size: int, dp_replicate: int) -> Tuple[int, int] | Tuple[None, None] staticmethod

Find batch size and accumulation steps that evenly divide dataset.

Parameters:

Name Type Description Default
dataset_size int

Total number of samples

 required
max_batch_size int

Maximum per-device batch size

 required
dp_replicate int

Data parallel replication factor

 required

Returns:

Type Description
Tuple[int, int] | Tuple[None, None]

(batch_size, accumulation_steps) or (None, None) if no valid size found

GPT

Bases: Module

embed(idx: jax.Array, deterministic: bool = False, **kwargs: Any) -> Any

Compute token and positional embeddings given inputs.

decode(x: jax.Array, padding_mask: Optional[jax.Array] = None, deterministic: bool = False, **kwargs: Any) -> Any

Compute decoded residual channels given embeddings.

measure(x: jax.Array, residual: jax.Array, depth: int, padding_mask: Optional[jax.Array]) -> None

Measure the residual update after one effective layer.

unembed(x: jax.Array) -> Any

Compute output distribution.

loss(logits: jax.Array, targets: jax.Array) -> jax.Array

Compute cross-entropy loss given logits and targets.

__call__(idx: jax.Array, targets: Optional[jax.Array] = None, padding_mask: Optional[jax.Array] = None, deterministic: bool = False, **kwargs: Any) -> Tuple[jax.Array, Optional[jax.Array]]

Parameters:

Name Type Description Default
idx Array

Input token indices of shape (B, T).

 required
targets Optional[Array]

Target token indices of shape (B, T). Use -1 to ignore positions.

 None
padding_mask Optional[Array]

Boolean tensor of shape (B, T). True for valid tokens, False for padding tokens.

 None
deterministic bool

If False, applies dropout.

 False

Returns:

Name Type Description
logits Array

Output logits of shape (B, T, vocab_size).
loss Optional[Array]

Cross-entropy loss if targets provided, else None.

Module

Bases: Module

sharding: ShardingPlan property

Logical mappings; models without a plan use replicated parameters.

flops(seq: int) -> float

Estimated forward + backward FLOPs for one sequence of length seq.

Composite modules require binding to existing variables so their actual setup children are available. Counts use dense matmul arithmetic (one multiply-add is two FLOPs, backward is twice forward), excluding optimizer updates, rematerialization and minor elementwise operations. Unsupported operations contribute zero; this is not an exact counter.

components() -> List[Type[Any]] classmethod

Return the types of constituent parts of this module.

Returns:

Name Type Description
Type List[Type[Any]]

A type or tuple of types representing the constituent parts.

gather() -> List[Type[Any]] classmethod

Depth-first search of all constituent parts of this module.

Returns:

Type Description
List[Type[Any]]

List[Type]: A list of all constituent part types.

Evaluate(spec: ExecutionSpec, base: Node | None = None)

Bases: Evaluator[GPT]

Evaluate initialized or restored GPT weights with a declared EVALUATION.

BackboneEvaluate(spec: ExecutionSpec, base: Node | None = None)

Bases: Evaluator[Module]

Evaluate HuggingFace weights or a native checkpoint based on that model.

job(key: str) -> Callable[[T], T]

Register a job class under the given key.

load_backbone() -> tuple[Any, Any]

Read the declared HuggingFace architecture and its host parameter tree.