dew

Dew: one registry, one objective, one trainer.

Naming anything the package exports fills the registries, so dew.models, dew.presets.EDM and dew.datasets["oxford_flowers102"] resolve after a bare import dew. The fill happens on that first name, not at import, so import dew.training stays inside the training layer and pulls in no modality, no encoder and no tracker backend, which is the layering rule tests/test_api_surface.py checks. Nothing here opens a JAX backend or loads an optional dependency; encoders, decoders and datasets fetch what they need when they are built.

objectives is not exported: dew.objectives is the package holding the Objective classes, and a registry cannot share its name. It is dew.registry.objectives.

Name	Summary
CFG	Interval-limited classifier-free guidance (Kynkaanniemi et al. Documented in dew.sampling.
Aux	Everything a loss returns besides its statistics. Documented in dew.objectives.base.
Checkpoints	Holds the checkpoints of one run, in one directory.
Condition	Names one conditioning input: its encoder, the batch field holding its tokens, and the raw datum for the unconditional branch. Documented in dew.inputs.
Dataset	Opens the batches a run trains and validates on. Documented in dew.data.
EMASpec	Say which leaves the EMA copy tracks, and how fast it follows them. Documented in dew.objectives.base.
Evaluation	Holds one evaluation event, with bounded hosted previews on rank zero. Documented in dew.training.
Field	A batch field and its per-example shape: Field("image", (128, 128, 3)). Documented in dew.inputs.
ImageGrid	Images in [-1, 1], [N, H, W, C], with the text each was conditioned on where there was any. Documented in dew.artifacts.
InputSpec	Names the sample field and the conditions, keyed by the model keyword each is passed under: {"textcontext": Condition(...)}. Documented in dew.inputs.
Layout	Says how a train state is placed on a mesh. Documented in dew.training.
LocalTracker	Writes synchronous reports in a tracking directory. Documented in dew.training.
MLflowTracker	An MLflow run in experiment, opened on the first value logged into it. Documented in dew.training.
MeshSpec	Says how many devices each sharding axis takes; data parallelism fills the rest. Documented in dew.training.
Objective	Define what is being learned: the parameters, the loss, what evaluation produces. Documented in dew.objectives.base.
Process	Pairs a schedule with what the model predicts on it and how the loss is weighted. Documented in dew.diffusion.process.
ProfileWindow	Asks for one profiler window per fit: steps steps traced into directory after warmup steps have run, so the trace holds the loop and not the compile. Documented in dew.training.
Representations	Encoder outputs [N, D] and the labels of the records they came from, for a probe to score. Documented in dew.artifacts.
Step	What the trainer tells an objective about the current step. Documented in dew.objectives.base.
TensorBoardTracker	A TensorBoard event file in directory, opened on the first value logged into it. Documented in dew.training.
TextSamples	Generated token rows, with optional decoded preview text and prompt. Documented in dew.artifacts.
TokenScores	Teacher-forced per-token losses [N, L] and the weight of each target, 1 where it counts and 0 where it is padding or a document’s first token. Documented in dew.artifacts.
Tracker	Documented in dew.training.
Trackers	Fan out without dropping reports; attempt every sink, raise the first failure. Documented in dew.training.
TrainState	Hold everything a run must checkpoint to resume where it stopped. Documented in dew.training.state.
Trainer	Runs an Objective: gradients, sharding, EMA, checkpoints, logging. Documented in dew.training.
VideoGrid	Clips in [-1, 1], [N, T, H, W, C]. Documented in dew.artifacts.
WandbTracker	A Weights & Biases run, opened on the first value logged into it. Documented in dew.training.
datasets	Documented in dew.registry.
encoders	Documented in dew.registry.
evaluate	Evaluate a finite coordinated prefix without an optimizer or tracker. Documented in dew.training.
metrics	Documented in dew.registry.
models	Documented in dew.registry.
pipeline	Load the inference task for source, its weights placed once. Documented in dew.inference.
presets	Documented in dew.registry.
profile	Configure native profiling; capture starts only on enter or start(). Documented in dew.telemetry.profile.
sample	steps points from T to 0: a solver step across each interval, then the model’s clean prediction at the last point. Documented in dew.sampling.
samplers	Documented in dew.registry.

Checkpoints

class source

class Checkpoints(
    directory: str,
    *,
    keep: int = 2,
    local_directory: str | None = None,
    local_every: int | None = None,
)

Holds the checkpoints of one run, in one directory.

Constructing one opens nothing; the orbax managers are created on first use. The directory keeps the latest keep steps, so a resume has something recent, plus the step with the lowest loss metric a save reported. A save without metrics can never become the best step.

local_directory names a path on every host’s own disk where the run keeps one more checkpoint, the latest, written every local_every steps by fit; each process writes the shards its devices hold under a directory of its own, so the same path serves a pod and a single host running several processes. latest is the newest step every process can read, local or persistent, and restore reads it from wherever it is. A local checkpoint restores onto the placement it was written with, since no process holds another process’s shards; the persistent checkpoint restores onto any mesh.

local_path: str

Return this process’s own local directory.

latest: int | None

Return the newest step a resume can read, local or persistent.

best: int | None

Return the step with the lowest reported loss, or None when no save carried one.

Checkpoints.path

def path(step: int) -> str

Checkpoints.source

def source(step: int) -> str

Return the directory restore reads step from: this process’s local one when the step is the local one every process holds, else the persistent one.

Checkpoints.save

def save(
    step: int,
    state: TrainState,
    saved: bytes | None,
    metrics: Mapping[str, float] | None = None,
    *,
    share: DataPartition | None = None,
) -> None

Write state under step, asynchronously.

Sharded arrays go straight to orbax: gathering them onto the host first would serialise the whole state through one process and undo the point of an async checkpointer. A stream reports its position as JSON bytes, which tensorstore has no dtype for; the raw bytes ride along as uint8 rows instead, one per process beside the data share it read, so a global position and a share’s offset are stored the same way and told apart on restore. A position without its share is refused, since no reader could be matched to it. A write that fails surfaces from wait, which is deliberately unguarded: a checkpoint that did not land is data loss.

Checkpoints.save_local

def save_local(
    step: int,
    state: TrainState,
    saved: bytes | None,
    *,
    share: DataPartition | None = None,
) -> None

Write state under step to this process’s local directory, asynchronously, in place of the local step before it. The placement rides along; a resume onto another one raises before reading shards from directories that do not hold them.

Checkpoints.stored

def stored(step: int | None = None) -> Variables

Return what the checkpoint at step holds, without reading its values.

step defaults to the latest. Each state field comes back as a shape/dtype tree, and an unset field as None.

Checkpoints.accumulation_template

def accumulation_template(step: int)

Return the persisted pending-array shapes, without reading their values.

Checkpoints.restore

def restore(
    template=None,
    step: int | None = None,
    *,
    share: DataPartition | None = None,
)

Restore the state at step and the data position of share.

template is a pytree of jax.ShapeDtypeStruct naming the state leaves to restore; a leaf’s sharding, when set, is where the array is placed, so a checkpoint written on one mesh restores onto whatever mesh this run is using. None restores every leaf as a host array.

A step that is the local one every process holds is read from the local directory, onto the placement it was written with; any other step from the persistent one. The data position comes back as the bytes the reader of share resumes from (read_position); without a share, as for a caller that reads weights and no data, it is None.

Checkpoints.wait

def wait() -> None

Block until pending async writes have landed on disk.

Saving is async so it stays off the training loop’s critical path; anything that reads a checkpoint back has to call this first.