dew.objectives.base

What the trainer is optimizing.

The trainer owns training mechanics: the mesh, the compiled step, EMA bookkeeping, checkpoints, logging. An Objective owns what is being learned: the parameter tree it initialises, the loss it computes from a batch, and what its evaluation produces. Swapping the objective swaps the research question without touching any of the mechanics.

An objective receives schedule and randomness through Step, and returns additive loss statistics with Aux reports. These values are JAX PyTrees.

Name	Summary
Task
Variables	A flax variables dict: the params collection plus any other collection the modules keep (moe, batch_stats, an objective’s frozen encoders).
Batch	One training example set, as the trainer and every objective read it.
Path
PathFilter	Selects leaves of a variables tree by the tuple of dict keys above them.
Initializer	An objective’s init as one value a JIT can take: a jax.tree_util.Partial, whose bound arguments are pytree children rather than closure cells, so a JIT that builds the initial state receives the held variables as arguments instead of compiling them in as constants.
Mean	Carry a scalar sum together with the mass it is averaged over.
mean_loss	Reduce a shared-denominator estimator, including empty support.
Loss
Effects
Step	What the trainer tells an objective about the current step.
Aux	Everything a loss returns besides its statistics.
Prediction	Hold what a token objective scored a batch with, for a teacher to compare.
everything
under	Build a filter that accepts the leaves below prefix.
select	Return the subtree of tree whose leaves keep accepts, nested the same way.
merge	Return tree with every leaf overlay holds replaced by the overlay’s.
FROZEN	The collection a partially trained run keeps its held weights under.
freeze	Move the params leaves trainable rejects under FROZEN.
thaw	Undo a frozen split, leaving one params collection again.
EMASpec	Say which leaves the EMA copy tracks, and how fast it follows them.
Objective	Define what is being learned: the parameters, the loss, what evaluation produces.
scalar_loss	Evaluate and reduce canonical statistics for direct JAX differentiation.
S
Metric	Reduce a validation pass to one scalar, on the host.
merge_totals	Add one batch’s (total, count) pair into a metric’s accumulator.
mean_of_totals	Divide a metric’s summed total by its summed count.

Task

attribute source

Variables

attribute source

A flax variables dict: the params collection plus any other collection the modules keep (moe, batch_stats, an objective’s frozen encoders).

Batch

attribute source

One training example set, as the trainer and every objective read it.

Declared once here and imported by dew.data.dataset, so a field added on one side is the same type on the other. The leaves are heterogeneous by measurement, not by omission: alongside the arrays a step consumes, a batch carries the prepared text inputs of a language run, the file paths a video corpus resolves lazily, the integer row counts a packer keeps and the record lists a mixture reads, so a narrower leaf type is untrue of what dew.data already builds.

Path

attribute source

PathFilter

attribute source

Selects leaves of a variables tree by the tuple of dict keys above them. One filter type serves the EMA selection, optax.multi_transform labels and frozen subtrees.

Initializer

attribute source

An objective’s init as one value a JIT can take: a jax.tree_util.Partial, whose bound arguments are pytree children rather than closure cells, so a JIT that builds the initial state receives the held variables as arguments instead of compiling them in as constants. A plain function is not one of these; jax.jit cannot take it as an argument.

Mean

dataclass source

class Mean()

Carry a scalar sum together with the mass it is averaged over.

The mass is nonnegative and does not depend on the parameters. Zero mass declares a zero numerator and no contribution.

mean_loss

function source

def mean_loss(stats: Mean) -> tuple[jax.Array, jax.Array]

Reduce a shared-denominator estimator, including empty support.

Loss

attribute source

Loss = TypeVar('Loss', default=Mean | jax.Array | float)

Effects

attribute source

Effects = TypeVar('Effects', default=None)

Step

dataclass source

class Step()

What the trainer tells an objective about the current step.

step is the count of accepted microbatches, which is what an objective’s own schedules index by. key is drawn fresh for every attempt, so a replayed microbatch draws the same randomness.

ema: Variables | None

The variables tree with the averaged leaves in place of the live ones, or None when the objective keeps no EMA.

Aux

dataclass source

class Aux(
    variables: Variables | None = None,
    qk_stats: Variables | None = None,
    effects: Effects | None = None,
)

Everything a loss returns besides its statistics.

metrics go to the tracker. variables replace nonparameter collections outright. effects are additive observations the optimizer applies once per commit rather than per microbatch.

variables: Variables | None

Complete nonparameter replacements from one accepted microbatch, such as BatchNorm statistics. The optimizer owns the params collection; deferred router-bias updates belong in effects instead.

qk_stats: Variables | None

The qk collection the attention layers sowed, for the optimizer’s QK-Clip: nested by module path, each attention layer holding max_logits as a one-tuple of an fp32 [rows, heads] array of per-head logit maxima, an MLA layer additionally holding qk_nope as an int scalar naming its nope width. Rows are the batch’s rows, microbatches concatenated under a pipeline. None when the loss never opened the collection, in which case the clip steps aside.

effects: Effects | None

Additive observations applied once on a supported optimizer commit.

Prediction

dataclass source

class Prediction()

Hold what a token objective scored a batch with, for a teacher to compare.

logits are the [B, S, vocab] fp32 scores the model’s forward produces, losses and weights the [B, S] per-position loss the objective sums and the weight it gives each position, and hidden the [B, S, D] states of the layers a caller asked for, in the order asked.

everything

function source

def everything(path: Path) -> bool

under

function source

def under(*prefix: str = ()) -> PathFilter

Build a filter that accepts the leaves below prefix.

As in under("params", "context_encoder").

select

function source

def select(tree: Variables, keep: PathFilter) -> Variables

Return the subtree of tree whose leaves keep accepts, nested the same way.

A branch that keeps no leaf is dropped, so the result is what the EMA stores and what merge puts back.

merge

function source

def merge(tree: Variables, overlay: Variables) -> Variables

Return tree with every leaf overlay holds replaced by the overlay’s.

FROZEN

attribute source

FROZEN = 'frozen'

The collection a partially trained run keeps its held weights under. The optimizer moves the params collection and nothing else, so what freeze leaves there is what trains; the rest rides beside it as state, and the model sees them merged by thaw.

freeze

function source

def freeze(variables: Variables, trainable: PathFilter) -> Variables

Move the params leaves trainable rejects under FROZEN.

Paths are full leaf paths, ("params", ...). A filter that keeps every leaf or none names nothing to split and is refused.

thaw

function source

def thaw(variables: Variables) -> Variables

Undo a frozen split, leaving one params collection again.

EMASpec

dataclass source

class EMASpec(decay: optax.Schedule, select: PathFilter = everything)

Say which leaves the EMA copy tracks, and how fast it follows them.

decay is a step-indexed schedule, since momentum ramps matter for some objectives (I-JEPA anneals 0.996 to 1.0). The step it reads is the count of completed optimizer updates.

Objective

class source

class Objective(ABC, Generic[Loss, Effects])

Define what is being learned: the parameters, the loss, what evaluation produces.

artifact: type | None

The artifact type evaluate returns, or None when it returns nothing.

inputs: InputSpec

Per-example shapes and dtypes the parameter tree is initialised from.

bank_sites: tuple[DecoderBank, ...]

Physical scanned stacks this objective evaluates, for an execution snapshot.

Each site names a decoder namespace below every variables collection and its StackView, as the model declares them. Objectives without a layer stack have only entry variables. The canonical variables and optimizer trees never adopt these banks.

initializer: Initializer

init as one value a JIT can take, with held arrays as arguments.

The trainer builds the initial state inside one JIT. A nullary function forces every concrete array its body reads to be captured as a compiled constant, which for a loaded checkpoint means the whole parameter tree is embedded in the executable: 2.2 GiB for a 0.6B model, a module too large for the compilation cache to store.

This is the one boundary where held variables cross into that JIT as data. Partial is a pytree whose bound arguments are children, so they arrive as JIT arguments however deeply init nests its own compilation, and the call always dispatches through public init.

Objective.held_variables

def held_variables() -> Variables | None

The arrays this objective starts from, or None when it draws them.

A continued-pretraining objective returns its checkpoint here; one that keeps a frozen tower beside the model it trains returns that. The trainer reads this once and hands the result back to init, so an objective never has to read its own held arrays inside a trace.

Objective.init

def init(key: jax.Array, variables: Variables | None = None) -> Variables

The whole variables tree, every collection, from one key. Pure, so the trainer traces it once for shapes and once for values.

variables is the held tree the caller supplies, which is how the trainer passes it as data; None means take it from this objective’s own held_variables. An objective that holds nothing ignores it.

Objective.loss

def loss(params: Variables, batch: Batch, step: Step) -> tuple[Loss, Aux[Effects]]

Additive loss statistics and the reports from one realized batch.

Mean declares a shared normalization mass. A plain scalar is one unit-mass term. Composite statistics are objective-owned Flax PyTrees; their leaves add across records before reduce_loss is evaluated.

Objective.reduce_loss

def reduce_loss(stats: Loss) -> tuple[jax.Array, jax.Array]

The objective value and whether its statistical support is active.

Objective.apply_effects

def apply_effects(variables: Variables, effects: Effects) -> Variables

Nonparameter replacements from accepted-window observations.

Objective.predict

def predict(
    params: Variables,
    batch: Batch,
    step: Step,
    *,
    train: bool,
    layers: Sequence[int] = (),
) -> tuple[Mean, Aux[Effects], Prediction]

The loss over batch as loss computes it, with the prediction behind it: the statistics, the reports, and the token logits with the weight of every position and the hidden states of layers.

The statistics are one Mean over the positions the weights count, so a distillation can mix in terms over the same mass. train gates dropout the way loss has it on; a frozen teacher scores with it off. Objectives that score no token logits raise.

Objective.evaluate

def evaluate(params: Variables, batch: Batch, step: Step) -> Artifacts | None

Scoring artifacts for every row of the coordinated global batch.

Every rank participates in numerical work outside the optimizer jit. No display sampling or decoding belongs here. Each scoring batch has a distinct key; step.ema holds the averaged weights.

Objective.pipeline

def pipeline(state: TrainState, *, ema: bool = True) -> Task

The trained model as its inference task over state’s weights.

Ordinary generative objectives require state.averaged when ema is True; False selects live parameters. Reference-policy objectives publish the trained policy, never their frozen loss reference. Arrays retain their placement. Objectives without a generation task raise.

Objective.preview

def preview(
    params: Variables,
    batch: Batch,
    step: Step,
    *,
    scored: Artifacts | None = None,
) -> Artifacts | None

One display per event, reusing first-batch scoring when available.

Called on every rank with a separate preview key. Before an internal collective, coordinate local setup and generation failures with agree_process_phase so every rank reaches the same boundary. Complete all gathers before root-only decoding. The trainer coordinates the hook’s final outcome before any subsequent collective.

scalar_loss

function source

def scalar_loss(
    objective: Objective[Loss, Effects],
    variables: Variables,
    batch: Batch,
    step: Step,
) -> tuple[jax.Array, Aux[Effects]]

Evaluate and reduce canonical statistics for direct JAX differentiation.

S

attribute source

S = TypeVar('S')

Metric

class source

class Metric(Protocol[S])

Reduce a validation pass to one scalar, on the host.

Statistics are merged as each batch arrives and finalized once.

The first contribution initializes a pass. State belongs to that pass alone; merge may update its owned buffers in place. Metrics must never perform process collectives or retain state between passes.

reads: type

The scoring artifact type this metric reads.

Metric.merge

def merge(accumulated: S, contribution: S, /) -> S

Combine a contribution with the pass-owned accumulator.

Metric.finalize

def finalize(accumulated: S, /) -> float

The completed pass’s scalar.

merge_totals

function source

def merge_totals(
    accumulated: tuple[float, float],
    contribution: tuple[float, float],
) -> tuple[float, float]

Add one batch’s (total, count) pair into a metric’s accumulator.

Every metric whose statistic is a sum over a count merges this way, so a pass over batches of different sizes still weighs by the count.

mean_of_totals

function source

def mean_of_totals(accumulated: tuple[float, float]) -> float

Divide a metric’s summed total by its summed count.