dew.training.state

Numerical state and retained accumulation records that cross jit.

Name	Summary
Accumulation	Hold the microbatches pooled so far, waiting for an optimizer commit.
Aux	Everything a loss returns besides its statistics. Documented in dew.objectives.base.
Step	What the trainer tells an objective about the current step. Documented in dew.objectives.base.
TrainState	Hold everything a run must checkpoint to resume where it stopped.
Variables	A flax variables dict: the params collection plus any other collection the modules keep (moe, batch_stats, an objective’s frozen encoders). Documented in dew.objectives.base.

Accumulation

dataclass source

class Accumulation()

Hold the microbatches pooled so far, waiting for an optimizer commit.

It keeps sums, not tapes: no forward residuals and no statistic Jacobians, so it survives a checkpoint. Statistics and effects retain their array leaves in canonical tree order, and the objective’s traced result supplies their PyTree structure.

Replay buffers have a leading window-slot dimension, then the original batch or mutable-collection dimensions. Only the collections Aux.variables rewrites need a per-record read snapshot.

TrainState

dataclass source

class TrainState(accumulation: Accumulation | None = None)

Hold everything a run must checkpoint to resume where it stopped.

Three clocks count separately. step counts attempts, and with the immutable root key it determines the next training draw. microstep counts accepted microbatches and indexes the objective’s schedules. updates counts committed updates and indexes the optimizer’s and the EMA’s schedules.

The scaler and the retained partial window are numerical state, and travel through the same checkpoint as the parameters.

averaged: Variables

The objective’s EMA leaves merged into the live variables.