API reference

ReactantNitro

A training framework for Lux models compiled with Reactant.

The central object is a Nitro. Nitro(e) runs setup only, so validate, evaluate and predict work without training; train! trains it. A run is one train! of a Nitro, which is what the run_* accessors are named for.

ASAM(rho; eta = 0.01f0, m = nothing)

Adaptive SAM (Kwon et al. 2021), which perturbs each parameter in proportion to its own size.

The gradient is taken at ps + ε with ε = rho * T² g / ‖T g‖ and T = |w| + eta, so the sharpness it measures is invariant to rescaling the weights. Leaves adaptive_perturb declines (by default the rank-1 leaves: biases, norm affines, scalar parameters) get T = 1, as under SAM. rho is relative, so typical values are about ten times SAM's, around 0.5 to 2 with SGD and smaller with Adam-family rules. eta is not schedulable; m is reserved, as for SAM.

Average(sel::Selector; weight = :uniform, state = nothing)

The parameters of the checkpoints sel picks from one run, averaged, as a weights source:

Nitro(e; weights = :latest => Average(Last(10)))Nitro(e; weights = "a1b2c3" => Average(Best(5; among = Last(20))))

weight is :uniform, the only weighting so far. Lux state is not averaged: it comes from one picked checkpoint, the best when sel ranks (a Best) and the latest otherwise, or as state = :best | :latest says. That is exact for a model without running statistics and an approximation for one with them, such as BatchNorm. The optimizer starts fresh, so restore_optimizer = true and resume refuse an Average. Averaging across runs (:all) is refused for now. The run must still hold every checkpoint sel names; keep them while training with Checkpointer(; keep), for example keep = (Best(3), Last(10)).

Best(k = 1; among = nothing)

The k best entries by the checkpointer's metric and mode. Only entries scored under that pair are ranked; under scope = :dir, every run's entries compete. among, a selector, ranks only what it picks: Best(5; among = Last(20)) is the best 5 of the latest 20. among is for loading, as in Average; keep refuses it.

CheckpointRecord

Everything a checkpoint holds: the snapshot passed to save_checkpoint!, plus format_version and framework_version, which the framework stamps.

Every value is a host value; setup moves them to the device when it loads a record.

Checkpointer(; keep = (Best(3), Last(1)), metric = nothing, mode = nothing, dir = nothing,
               name = nothing, scope = :run)

The default checkpointer. After each epoch it writes a full CheckpointRecord into run_dir if some selector in keep would retain it, then keeps the union of what the selectors pick and deletes the rest. The default, (Best(3), Last(1)), keeps the best 3 by metric (lower is better under mode = :min, higher under :max) plus the latest, whatever it scored. checkpointer = nothing disables checkpointing.

Checkpointer(; keep = (Best(3), Last(10)))   # the best 3 and the latest 10, for averagingCheckpointer(; keep = Last(5))               # a rolling windowCheckpointer(; keep = Best(3))               # the best 3 only

metric and mode default to nothing, which follows the run's monitor; give both to rank by something else. Setup fills a following checkpointer in, every time it builds a Nitro. A checkpointer used without setup (in average_checkpoints, say) ranks by the metric and mode the run's newest scored entry was written under. dir and name left as nothing are filled in at setup: dir with run_dir, and name, the filename function name(; epoch, step, metric, score), with the experiment's checkpoint_filename.

A manifest, manifest.jld2, lists each retained file with its epoch, score, stop_reason, and the run, metric and mode that wrote it, so lookups and resume never open a record; see read_manifest. The latest checkpoint is found through it rather than a latest symlink, which would dangle once rotation deleted its target.

Without a Last, an epoch no selector admits is not written at all, rather than written and deleted. The final stop_reason rewrite follows the same rule, so the state training ended in is saved only if a selector keeps the last epoch, and weights = :latest => :latest and resume = :auto resolve to the latest kept epoch, not to where training stopped. A crash loses every epoch since then. Setup warns about this for a handle that trains, and refuses keep with only Best selectors for a training run with no val split, whose epochs have no score.

scope = :run, the default, applies retention to this run's checkpoints only. A directory often holds several runs (the default run_dir is the experiment's type name), so setup gives each run an id, the field run, and never ranks, rotates, deletes or overwrites another run's checkpoints. A fresh run gets a new id. A resume from a run's latest record continues that run's id; a resume from any other record starts a new id and records its source path in parent. Setup sets both fields for every Nitro it builds, so a checkpointer reused across handles belongs to the most recent one. A checkpointer used without setup has run = nothing, the same unnamed run as manifest entries written before runs were recorded.

scope = :dir saves space: Best ranks the whole directory, across every run, and a write deletes every other entry scored the same way, including other runs'; Last and the other selectors keep this run's own entries. Setup warns a training handle when the directory already holds checkpoints this may delete.

Only entries scored under the current metric and mode are ranked. Entries written under another pair, by an earlier run or before a resume changed the selection, are never ranked or deleted, under either scope. Setup and :best lookups warn when the entries they consider are mixed this way. An entry with no recorded metric predates the field and counts as matching. Under either scope, a filename that would overwrite a file this run does not own gets -run-<id> inserted before .jld2.

ckpt.k (the largest Best's k, else 0) and ckpt.keep_latest (whether keep has a Last) remain readable for code written against TopKCheckpointer, and are deprecated.

Writing a checkpoint.

Compiling <: Phase

Supertype of the four compile phases, which normally take minutes. Test phase isa Compiling rather than a leaf type, so a monitor keeps working when a compile phase is added.

Decay(lambda)
Decay(lambda, anchor)
Decay(lambda, anchor, no_decay_mask)

Decoupled weight decay toward an anchor: apply! adds no_decay_mask * lambda * (x - anchor) to the update, or no_decay_mask * lambda * x when anchor === nothing. The anchor comes from the group's decay_anchor:

A group whose lambda is zero gets no Decay in its chain. anchor === nothing is resolved by type, so plain decay emits no subtraction. Decay must be the last rule in a chain, after the base rule as in AdamW, so lambda arrives pre-multiplied by the group's learning rate; any other position is a setup error. The anchor is a field rather than state captured at init, which on a resume would capture the restored weights. There is no WeightDecay alias, since Optimisers.WeightDecay already uses that name.

Device{T}

Declaration-time marker for an @experiment field; the macro consumes it and records the field in device_fields. A Device field is converted to device residency at setup and reaches the traced step as an input, so e.aux_weight in a hook is the device value. The marker places the value on the device, keeps it out of the compile cache key (a traced input cannot change the graph), and makes it schedulable (a constant and a schedule use the same slot). Mark what you intend to change or schedule without recompiling; a Device threshold of zero still emits ops that a host 0f0 would fold away.

@experiment struct MyExp    "Weight of the auxiliary heatmap loss relative to the primary term."    aux_weight::Device{Float32} = 0.25f0end

See also Host, the default, and GraphConst, which bakes.

Done <: Terminal

The run finished: it reached max_epochs, stopped early, or stopped on request_stop!. The handle's stop_reason says which.

EMA(beta; warmup = true, every = 1)

An exponential moving average of the weights, θ_ema ← β θ_ema + (1 - β) θ, updated after every every-th optimizer step. warmup uses min(beta, (1 + t) / (10 + t)) at optimizer step t, so the average is not dragged by the initial weights for 1 / (1 - beta) steps. Set it with the ema accessor or the Nitro keyword:

ReactantNitro.ema(::MyExp) = EMA(0.999)

The average starts as a copy of the weights training starts from, is checkpointed in the record's ema field and restored by a resume. Under eval_weights = :ema (the default with an EMA), validation, checkpoint selection, predict, restore_best and export use it; training and parameters use the raw weights. Nitro(e; weights = ..., load_ema = true) loads a record's average in place of its weights. The average has the weights' element type.

EarlyStopping(; metric = nothing, mode = nothing, patience = 5, min_delta = 1f-4)

Stop training when a validation metric stops improving. Pass it as the early_stop keyword or accessor; the default nothing never stops early.

metric and mode default to nothing, which follows the run's monitor; give both to watch something else. patience counts epochs without improvement, and an epoch improves when it beats the best so far by more than min_delta (absolute, as in Keras and Lightning). Setup requires a val split and checks mode, patience and, for an experiment without metrics, that metric is :loss. With a custom metrics the key cannot be checked before it runs, so a misspelled one raises after the first epoch's validation.

train!(e; early_stop = EarlyStopping(; patience = 5))                              # the monitortrain!(e; early_stop = EarlyStopping(; metric = :loss, mode = :min, patience = 5))

The stopping epoch still validates and checkpoints, then train! returns through Done with stop_reason = :early_stop, which the final checkpoint record stores. request_stop! stops at the same point on demand, with stop_reason = :requested. Setup copies the instance with the monitor filled in and fresh best and wait counters, so two handles never share them and the instance you passed is not mutated; the handle's copy is nitro.early_stop. The counters are stored in each checkpoint, and a resume that watches the same metric and mode continues from them.

Epochs(epochs)

This run's entries whose epoch is in epochs, a range or a collection.

Tracing and compiling the eval-mode forward / metrics programs.

EvalStepping <: Stepping

Running the validation, testing, or inference loop.

Every(n)

This run's entries whose epoch is a multiple of n.

ReactantNitro.ExportBackend

The supertype of an export target. A backend implements one method, write_export. The methods for ReactantServerBundle and TFSavedModel live in package extensions, so this package depends on no artifact format.

ExportCompiling <: Compiling

Tracing and compiling the export program. Set once per export_model call, around the backend's write_export, which compiles outside the compile cache. The previous phase is restored when the bundle is written or the call fails.

ExportSpec(name; from = nothing)
ExportSpec(name, dtype, shape; batch_axis = length(shape), axis_letters = nothing, from = nothing)

One tensor's declaration for export, which a backend translates into its own format. shape is the Julia shape and batch_axis a 1-based Julia axis. Which fields are required depends on the hook:

The executable side is traced, so its shapes are known and batch-last is enforced. The client side is what a model.jl the framework never runs produces or receives, so it must be declared in full. axis_letters names the non-batch axes in shape order, e.g. ['w', 'h']; a -1 in shape marks a variable non-batch axis. from gives an output a source other than its own name, and is accepted only on export_outputs.

The run raised an exception.

Tracing and compiling the gradient program.

GradRule

The supertype of what grad_rule returns.

GraphConst{T}

Declaration-time marker for an @experiment field. A GraphConst field bakes as a trace-time constant and enters the compile cache key: the tracer sees a real Int with ordinary control flow, and a changed value recompiles. Use it for structure: layer counts, widths, variant selectors. A field left unmarked by mistake is Host, and traced code that reads it raises a StrippedHost error instead of compiling a wrong program.

@experiment struct MyExp    "Number of decoder blocks. Structural: changes the compiled graph."    n_layers::GraphConst{Int} = 4end

Host{T}

Declaration-time marker for an @experiment field, and the default for an unmarked one: max_epochs::Int = 40 means max_epochs::Host{Int} = 40. A Host field is never converted, never hashed into the compile cache key, and never visible to the tracer. This keeps driver knobs out of the key, and keeps dataset-sized state away from Reactant and Enzyme, which walk the whole Const(e) argument element by element at every compile. compile_view replaces every Host field with a StrippedHost sentinel at every trace site.

See also Device and GraphConst.

JSONLogger(path) -> JSONLogger
JSONLogger() -> JSONLogger

The default logger: writes JSON Lines to path, or, when constructed without one, to metrics.jsonl in the run's run_dir. logger(e) returns a pathless one when the experiment declares no logger, so a plain Nitro(e) writes runs/<Exp>/metrics.jsonl beside its checkpoints. Each line is one object with a "type" key (params, tags, metrics, other, confusion or finish); metric lines also carry context, epoch and step:

{"seed":42,"max_epochs":40,"width":128,"type":"params"}
{"loss":0.31,"type":"metrics","context":"train","epoch":1,"step":4}
{"loss":0.29,"type":"metrics","context":"validate","epoch":1,"step":160}
{"data_wait_frac":0.012,"type":"metrics","context":"data","epoch":1,"step":160}
{"type":"other","key":"binding_report","value":"..."}
{"type":"finish","status":"completed"}

The file is opened in append mode and flushed after every line, so a resumed run extends the same file. Pass logger = nothing to log nothing.

Last(n = 1)

This run's latest n entries: by epoch under scope = :run, by write time under scope = :dir.

Nitro(e; kwargs...) -> Nitro

The state of a run: model, weights, optimizer, data and counters. Nitro(e) runs the setup sequence and nothing else, so validate, evaluate and predict work without training in the process. train!(e; kwargs...) is shorthand for train!(Nitro(e; kwargs...)), so every run keyword is a Nitro keyword.

nitro = Nitro(e)                                  # fresh weights from build_modelnitro = Nitro(e; weights = :latest => :best)      # best weights of the newest run in run_dirnitro = Nitro(e; data = (; test = loader))        # supply data directly, skip build_datanitro = Nitro(e; compile_options = (; cudnn_hlo_optimize = true,    # Reactant compile keywords,    xla_debug_options = (; xla_gpu_exhaustive_tiling_search = true))) # XLA flags among them

compile_options applies to every program the handle compiles and is part of each program's compile-cache key; see ReactantNitro.check_compile_options.

Read it through accessors: experiment, parameters, states, run_dir, current_step, current_epoch, phase, history and logger_info; request_stop! is the one write. It holds no Enzyme shadow dps, which exists only inside the gradient program.

Nitro(e; seed, resume, run_dir, data, n_devs, weights, accum, max_epochs, schedules,
         gradient_clip_norm, check_divergence, autolog_schedules, log_every_n_steps, logger,
         checkpointer, early_stop, monitor, eval_loss, ema, run_ref, w0, restore_optimizer, restore_best,
         eval_weights, load_ema, compile_options) -> Nitro

Build a handle for experiment e by running the setup sequence below. Nothing is trained or compiled.

Sixteen keywords default to the accessor of the same name: seed, run_dir, n_devs, accum, max_epochs, schedules, gradient_clip_norm, check_divergence, autolog_schedules, log_every_n_steps, logger, checkpointer, early_stop, monitor, eval_loss and ema. Passing one replaces the accessor's value for this run; omitting it calls the accessor, before step 1. data, weights, resume and run_ref have no accessor, since each names a fact about this construction: data supplies the collection and skips build_data, and run_ref, a Ref, receives the handle when construction finishes. compile_options passes Reactant compile keywords to every program the handle compiles (see ReactantNitro.check_compile_options).

The fresh sequence

Steps 9 and 11 run after step 10. Each program is traced and compiled, against compile_view(e), by the first verb that runs it.

build_data and derive see the pre-conversion experiment, where Device fields hold host values; every other hook sees the post-conversion one.

Resume adds steps: before 1, locate and load the checkpoint; at 2, take seed from the record, with a warning if it differs, since a seed sweep sharing one run_dir would otherwise become N identical runs without notice; after 4, check config compatibility; after 7, restore ps and st, keeping the fresh w0 capture to verify against the record's anchor_checksum; after 8, refuse a flat permutation or anchor checksum that differs; at 9, restore opt_state and re-normalize it to device residency, without which a resumed run trains at the wrong point of its optimizer's bias correction with no error. step and epoch come from the record.

Checkpoint sources. weights and resume name a checkpoint the same way: a path, or run => checkpoint, resolved to a file through the manifest (resolve_source).

Best(1) and Last(1) are the same as :best and :latest; any other Selector may pick several checkpoints, so a source takes it only inside Average.

Bare :best and :latest are refused, naming the pair to use (weights = :latest => :best), as is :all => :latest, which is :latest => :latest. A source that cannot be found raises, naming the directory and the runs it holds. :best selects among entries scored under the current metric and mode only, warns when the run also holds entries under another, and raises when it holds none under the current one, listing the ones it has.

resume is always exact continuation: parameters, layer state, optimizer state, step, epoch and schedule position, under the compatibility checks. It is false (the default), a path, run => :latest | :best with any run above including other, or :auto, which is exactly :latest => :latest and logs what it resolved to, except that a directory with nothing to resume starts a fresh run with a warning (:latest => :latest raises there). One run_dir can hold several runs, each with its own id in the manifest; see Checkpointer. Resuming a run's newest record continues that run: its id is adopted and later checkpoints join it. Resuming anything else (an earlier record such as :best, :all => :best, or a file this directory's manifest does not list) branches into a new run, so the original run's later epochs are never overwritten, and logs that it does. Under a checkpointer with no Last in keep (what keep_latest = false builds) a run's latest record is its most recent kept epoch rather than where it stopped, and the epochs after it are trained again. resume = other and other => :current, a continuation from the handle's in-memory state, are not supported yet; weights = other takes the parameters.

weights names where the initial parameters come from, and takes no training position: no step, no epoch, no schedule position, and no optimizer state unless restore_optimizer asks. A handle built this way that trains is a new run in run_dir, and a directory already holding checkpoints keeps them (under the default retention scope = :run).

• weights = path or run => :best | :latest restores ps and st from a checkpoint record, and the derived values with them, so derive is skipped and the construction works with no training data. An older record holds no list of derived keys, so for it derive runs and the record's Device values override what it returns. Reloading a finished run's best weights needs no file name:

  n = Nitro(MyExp, :baseline; run_dir = "runs/my_run", weights = :latest => :best, logger = nothing)

• weights = run => Average(sel) restores the average of the run's checkpoints sel picks, with st and the derived values from one of them (see Average), and records what it averaged in averaged_from. average_checkpoints writes the same record to a file.

• weights = other::Nitro, or other => :current, is a warm start from a handle in this process: ps and st are taken after step 7, through host memory, and otherwise the fresh sequence runs. The trees must match leaf for leaf.

restore_optimizer = true also restores the optimizer state of a weights record (a path or any run => :best | :latest), normalized to device residency as a resume normalizes it and checked against the optimizer this run builds (rule and group layout; a mismatch is refused), while step, epoch and the schedule start fresh: a new run with a new schedule and max_epochs. It refuses a source that is not a record (none, a bare Nitro, other => :current: restoring a live handle's optimizer is not supported yet), refuses alongside resume, which restores the optimizer anyway, and refuses on a handle with no train split, which has no optimizer to restore into.

restore_best = true ends every train! that does not fail (completion, early stop, a requested stop) by loading this run's own best checkpoint, under the current metric, into ps and st; step, epoch, history and the last metrics stay as training ended. checkpoint_source becomes that file, so export provenance names it, and show says the handle holds it. With no scored checkpoint it warns and keeps the final weights. A handle restored this way refuses a later train!, and the error shows the three constructions that continue from the best checkpoint.

restore_best = Average(sel) ends train! holding the average of the run's own checkpoints that sel picks instead, exactly as weights = run => Average(sel) would load them. Setup refuses a checkpointer whose keep cannot hold that selection; a run that ends before it has enough (an early stop, say) warns and keeps its final weights.

eval_weights is :ema (the default with an EMA) or :raw (the default without one), and says which weights validation, checkpoint selection, predict and export read. load_ema = true loads the weights source's EMA in place of its weights: a record's, an average of records' (run => Average(sel)), or another handle's. A source without one is refused.

w0 says what decay_anchor = :w0 anchors to: :build_model (the default), :weights (the transferred parameters, for L2-SP fine-tuning), or a parameter tree.

checkpoint = path is the deprecated spelling of weights = path. It still works, warns once, and will be removed in a later release.

Without a train split steps 9 and 11 are skipped. Routing falls back to the first batch of whichever split exists; with no split at all it is deferred to the first predict call.

Progress. The sequence is reported as one "setup" stretch of unknown length, with the current step as its phase (setup [building data]), closing as setup done. See progress_reporter!.

Nitro(E::Type, preset::Symbol; kwargs...) -> Nitro

Build the experiment from from_preset(E, preset) and record which named recipe the run claimed, so the name reaches the checkpoint record and log_params!.

n = Nitro(MyExp, :baseline; max_epochs = 40, aug_rotate_deg = 9.0)#                              ^ run keyword    ^ field, so it overrides the recipe

Keywords are split by one rule: a Nitro keyword goes to Nitro, and anything else must be a field of E and goes to the recipe. A name that is both, such as max_epochs or seed, is the run keyword. A name that is neither is an error listing both sets. data is refused in this form, since the collection would come from a different instance than the one built here; use the longer spelling below.

This is shorthand for Nitro(from_preset(MyExp, :baseline; ...); preset = :baseline), which still works but names the preset twice, and nothing checks that the two agree.

ReactantNitro.NoPrefetch(source)

Opt a split out of prefetch, so its host data path runs inline on the training task. This is the only way to do so; device_batches = 0 is an error, so the choice is visible in the code rather than looking like tuning. iterate passes through to the source, as PrefetchIterator's does.

Tracing and compiling the optimizer program.

PerturbRule <: GradRule

Rules that take the gradient at perturbed parameters, ps + ε, where ε comes from the unperturbed gradient. A subtype defines perturbation; the driver runs the gradient program twice per micro-batch, at ps and then at ps + ε, and accumulates the second gradient. The reported loss, the layer state and the train_metrics outputs are the first pass's.

Phase

Root of the run-phase tree. The framework publishes phase transitions to the monitors registered with register_phase_monitor!, so a heartbeat, watchdog or dashboard can be written outside the framework. Phase durations differ by orders of magnitude (a compile takes minutes, a step milliseconds), so a watchdog needs to know the phase; the framework ships no timeouts.

Plateau(inner; metric = nothing, mode = nothing, factor = 0.1, patience = 10, cooldown = 0,
        threshold = 1f-4, threshold_mode = :rel, min_scale = 0)

A schedule factory that multiplies inner by a scale, and reduces the scale when the validation metric stops improving. inner is any schedule factory or a bare Number; a bare number is PyTorch's ReduceLROnPlateau. metric and mode default to nothing, which follows the run's monitor; give both to watch something else. It works on every schedule key: opt (bare and group paths), device, grad and bind.

schedules(e::MyExp) = (;    eta = Plateau(total -> OneCycle(total, 1f-3); patience = 3),)

The value at step t is inner(t) * scale. The scale starts at 1 and is updated once per epoch, after validation, with PyTorch's rule and defaults:

• an epoch improves when the metric beats the best so far by more than threshold, relative (best * (1 - threshold) for :min) under threshold_mode = :rel, or absolute (best - threshold) under :abs
• after more than patience epochs in a row without improvement, scale = max(scale * factor, min_scale), and the next cooldown epochs count no bad epochs

So patience = 2 reduces at the third bad epoch, as in PyTorch, and patience = 0 reduces at every bad one. A reduction applies from the next epoch's first step.

Each value is a pure function of the step and the run's validation history. The counters are stored in each checkpoint, so a resume continues the same sequence; a weights load starts fresh. Each key gets its own counters. The train line logs the scaled value as usual, the validation line logs schedule/plateau/<key>, and each reduction is reported with @info.

Setup requires a val split, mode in (:min, :max), 0 < factor < 1, patience >= 0, cooldown >= 0, and, for an experiment without metrics, metric = :loss. With a custom metrics, the key cannot be checked before it runs, so a misspelled one raises after the first epoch's validation.

ReactantNitro.PrefetchIterator(source; workers = default_prefetch_workers(),
                               device_batches = 1, host_batches = 2 * workers, ordered = true)

Prefetch settings for a split: workers producer tasks build host batches, one transfer task copies them to the device, and a Channel of device_batches (host, device) pairs feeds the loop.

The framework wraps every split in one of these at the defaults (see auto_prefetch). Construct one only to change the defaults; NoPrefetch opts a split out.

• workers sets throughput. Buffering only adds lookahead: a producer slower than the device starves it at any device_batches. More than one worker needs batch_at and begin_epoch! on the source; without them the split gets one producer, and setup warns for the training split.
• device_batches bounds batches staged on the device. Resident device memory is (device_batches + 2) x batch (the channel, the transfer task's hand, the consumer's hand), hence the default of 1.
• host_batches bounds batches on the host at once, across the workers, the hand-off channel and the reorder buffer. It must be at least workers, or ordered delivery can deadlock waiting on a batch that was never handed out.
• ordered = true emits in source order through a reorder buffer, so a fixed seed reproduces bitwise at any workers. ordered = false emits batches as they finish, so the epoch's accumulation groups change and runs no longer reproduce bitwise; every batch is still seen exactly once. Applies only to the fan-out; one producer always emits in order.

The transfer runs on one task because device memory cannot be bounded by a semaphore: XLA runs asynchronously and the executable holds its inputs, so there is no point at which the consumer is known to be done. Device buffers are freed explicitly, not left to the GC.

Batching, shuffling and splitting are not provided; use MLUtils.jl.

The wrapper holds only the source and its settings, since build_data runs before batch routing and device placement exist; the training loop builds the pipeline. Base.iterate therefore passes through to the source with no task and no channel, so iterating one by hand cannot leak a producer, and render and predict(nitro, loader) are unaffected by the wrapping. Eval splits stream through eval_stream, which pads a short final batch before the transfer.

ReactantNitro.ReactantServerBundle()

The StableHLO bundle format served by ReactantServer. The type is defined here so it can be named; its write_export method lives in the ReactantServerExport extension, so an export needs using ReactantServerExport.

Repl <: Phase

The caller has control and no framework work is in flight, as after any public entry point returns. It lets a monitor tell a finished run from a process stuck in teardown.

Published but never recorded: Repl describes the process, not the run, so phase still reports how the run ended (Done or Failed). Only the outermost entry point publishes it, tracked by a process-level depth counter (work_in_flight), so train! calling validate each epoch does not announce an idle session mid-run.

SAM(rho; m = nothing)

Sharpness-aware minimization (Foret et al. 2021): the gradient taken at ps + rho * g / ‖g‖.

rho is a host number; the framework places it on the device, so changing it never recompiles. Schedule it as rho = ..., as grad = (; rho = ...) when another namespace has a rho, or by hand through bind and grad_rule(e, hp). ε is computed per accumulation micro-batch and excludes the leaves no_perturb excludes, none by default. The cost is two forward and backward passes per micro-batch. m, a per-chunk sample count for independent m-sharpness, is reserved and not supported yet.

Selector

Which entries of a run's manifest a rule picks. Checkpointer(; keep) retains the union of its selectors; a weights source picks with one. Best, Last, Epochs and Every are built in. A custom selector defines select and, optionally, admits.

SingleGrad()

The default grad_rule: one forward and backward pass per micro-batch.

Starting <: Phase

Work has started and has not yet reached a compile or a step. Published when Nitro(e) begins setup and when an outermost entry point starts; a new handle's phase is Starting until its first transition.

Stepping <: Phase

Supertype of the two per-batch phases. A run alternates between them.

ReactantNitro.TFSavedModel(; platform = "CUDA", stablehlo_version = v"1.5.0",
                           call_module_version = 9)

A TensorFlow SavedModel for TF Serving and the platforms built on it. Each traced program is one XlaCallModule op, the form JAX's native serialization produces, and weights are non-trainable variables shared by every signature. Its write_export lives in the PythonCall extension, which drives TensorFlow to write the files, so an export needs using PythonCall and an importable tensorflow.

There is one signature per entry of batch_sizes, serving_b{N}, with serving_default the first. Tensors are batch-first: the Julia shape (W, H, C, N) is the TF shape [N, C, H, W]. model.jl cannot ship, so export_client_outputs is refused and any postprocess is dropped with a warning. Provenance is written to assets.extra/.

platform is the single platform the module runs on ("CUDA", "CPU", "ROCM", "TPU"). The defaults load on TF 2.17 and later; TF 2.15 needs stablehlo_version = v"0.14.0" and call_module_version = 8. A module using an op newer than stablehlo_version fails to export.

Terminal <: Phase

Supertype of the two end states. A monitor tests p isa Terminal to know when to close.

Running training steps.

TopKCheckpointer(; k = 3, metric = nothing, mode = nothing, dir = nothing, name = nothing,
                   keep_latest = true, scope = :run) -> Checkpointer

Deprecated: use Checkpointer(; keep = (Best(k), Last(1))), or keep = Best(k) for keep_latest = false. Builds exactly that Checkpointer, so the files and the manifest are the same as before. Warns once per session; supported until at least 1.0.0. metric and mode follow monitor as Checkpointer's do.

accum(e) -> Int

Micro-batches per optimizer step. Default _field(e, :accum, 1); also a Nitro keyword. It bakes into the gradient program as the constant inv_n = 1/accum, so changing it recompiles that program, and a GraphConst field of this name is allowed. The train split's batch count must divide by it.

adaptive_perturb(e, keypath::Tuple, param) -> Bool

Whether ASAM scales this leaf's perturbation by its own size, |w| + eta, rather than leaving it unscaled as SAM does. Default: rank-2 and higher leaves (weight matrices and kernels); biases, norm affines and 1-element scalar parameters are unscaled. Runs on the host once, at construction.

admits(sel::Selector, pool, ckpt, epoch, score) -> Bool

Whether sel would retain a write of epoch scoring score, decided before the record is built. An epoch no selector admits is not written. The default is true, which always writes and leaves the decision to select.

autolog_schedules(e) -> Bool

Whether every varying scalar schedule value is logged on each train line (see log_every_n_steps). Default _field(e, :autolog_schedules, true); also a Nitro keyword. A Device field logs as schedule/<field> and a rule value as schedule/<group>/<field>, after the per-group scaling (see Schedules); constant schedules are left out. Host-side only: the values are the ones the step already computed.

average_checkpoints(run_dir, run => Average(sel); checkpointer = Checkpointer(), save = nothing)
    -> CheckpointRecord

The averaged record weights = run => Average(sel) would load, returned on the host, and written to save when one is given. The written file is an ordinary record: weights = path loads it, and its averaged_from field (shown by checkpoint_info and recorded in an export's provenance) names the files, epochs and weights it averages. It is not added to the manifest, so retention never ranks or deletes it. save refuses an existing file.

average_checkpoints("runs/MyExp", :latest => Average(Last(10)); save = "runs/MyExp/avg-last10.jld2")

backend(lgr) -> Any

The native backend object, for calls the contract does not cover (images, artifacts, curves). The default is identity, since a logger is usually the backend object itself; a wrapper returns what it wraps. Phase monitors reach the logger as info.logger.

backend(w::MyWrapper) = w.exp

backward(f, ps_sub, consts...) -> (loss, grads)

One reverse-mode pass over an objective with respect to one parameter subtree, for use inside a train_step closure. f(ps_sub, consts...) returns the scalar loss; grads matches ps_sub, and everything in consts... is Const, so no gradient flows into it.

Every traced value the objective reads must be an argument, never a closure capture: a captured parameter subtree gets a zero gradient with no error, while the loss is still correct. A struct holding no arrays, such as the model, is safe to capture.

For a GAN generator step, pass the discriminator's parameters in consts...: the generator's gradient flows through the discriminator as a function of the fake data and never into its parameters. The gradient buffer is allocated per call, as in the gradient program.

ReactantNitro.batch_at(source, i::Integer) -> batch
ReactantNitro.batch_at(source, i::Integer, plan) -> batch

Produce batch i of the current epoch. Optional; with begin_epoch! it opts a source into the multi-producer fan-out, which calls it from many concurrent tasks, each i independently. A source that stores its own plan defines the two-argument form. An immutable source has begin_epoch! return the plan and defines the three-argument form; leave plan untyped, since fanout_capable detects the method by asking whether it accepts Any.

i is a BATCH index in 1:length(source), not a sample offset:

ReactantNitro.batch_at(dl::MyLoader, k::Integer) =    _build_batch(dl, 1 + (k - 1) * dl.batch_size)     # RIGHTReactantNitro.batch_at(dl::MyLoader, k::Integer) = _build_batch(dl, k)   # WRONG, and it TRAINS

The wrong version trains each epoch on overlapping windows of the first few hundred samples, and the loss still falls; check_batch_at catches it. An implementation must be thread-safe with respect to shared state (keep per-task scratch in a TaskLocalValue), depend only on i and the epoch's plan (re-planning belongs in begin_epoch!), and never return nothing for an i in 1:length(source).

ReactantNitro.begin_epoch!(source) -> plan

Plan the next epoch (shuffle, draw augmentation seeds). Optional; with batch_at it opts a source into the multi-producer fan-out. The fan-out calls it once per epoch, before requesting any batch. Return nothing if the source stores its own plan, and the framework calls the two-argument batch_at; otherwise return the plan, and every three-argument batch_at call that epoch receives it.

The fan-out never calls Base.iterate, and the other paths never call begin_epoch!, so a source should re-plan here and have its iterate call this method, as below. A source that re-planned in iterate and defined only batch_at would reuse epoch 1's plan for every epoch, so both methods are required; a training split with only batch_at gets one producer and a setup warning.

ReactantNitro.begin_epoch!(::MySource) = nothing            # no re-plan neededReactantNitro.begin_epoch!(d::MyLoader) = shuffled_index(d.rng, d.n)   # a declaration's planBase.iterate(s::MySource) = (begin_epoch!(s); _first_batch(s))         # one implementation

length(source) is read after this returns, so a plan may change the batch count. Setup's schema probe goes through Base.iterate, so a source that counts its own epochs sees N + 1 calls across N epochs, and a loader that plans eagerly in build_data needs a flag saying the first plan is already drawn.

build_data(e, dist) -> NamedTuple

Required. Return the run's named data collection, (; train, val) or (; train, val, test).

A data source is anything iterable that yields concrete NamedTuples of host arrays and supports length, which counts batches: a Vector of batches, an MLUtils.DataLoader, or a bespoke sampler. The framework transfers each batch, shards it, and owns prefetching. Prefer MLUtils.DataLoader: its settings are checked at setup from its own fields, the MLUtils extension fans it out across every thread, and batching, shuffling and collation are its job. A bespoke sampler puts those back on you.

Three requirements: the source is restartable (setup draws one batch for the schema and the loop iterates from scratch, so a bare Channel loses a batch); length is read once at setup to fix the schedule horizon; and the train split drops its partial final batch (partial = false, drop_last = true) and its batch count divides by accum. Eval splits should not drop theirs, since the framework pads and slices, and testmode makes a standard model per-sample along the batch axis.

There is no prepare_epoch! hook; a loader re-plans in its own iteration initialization. dist is nothing today; write it untyped. build_data sees the pre-conversion experiment, and is skipped when Nitro(e; data = ...) supplies the collection.

build_model(e, rng) -> (model, ps, st)

Required. Return a Lux model with its parameters and state. Load pretrained weights here: the framework captures w0 as soon as build_model returns, so weights loaded later would leave decay_anchor = :w0 anchored to the random init. It sees the post-conversion experiment. For a model with variant branches, lift the selector to a type with dispatch_variant so the branches fold at trace time:

build_model(e, rng) = _build_model(e, dispatch_variant(e, :model_kind), rng)_build_model(e, ::Val{:bit_r50}, rng) = ...

check_divergence(e) -> Bool

Whether a non-finite training loss stops the run. Default _field(e, :check_divergence, true); also a Nitro keyword. When on, a NaN or Inf loss raises, naming the step and epoch (see train!). The automatic loop reads each micro-batch's loss back one micro-batch late, so the error comes after at most one more micro-batch has been dispatched. When off, the loss is still read back; a non-finite value is dropped from the logged train line like any other, the epoch's mean train loss in history becomes non-finite, and training continues from the diverged state. Host-side only, so changing it recompiles nothing.

checkpoint_filename(e; epoch, step, metric, score) -> String

The filename of a checkpoint; a hook a model package may extend.

epoch-0006-step-4500-loss=0.0416831.jld2epoch-0006-step-4500.jld2                          # a run with no `val` split

The epoch comes first, zero-padded, so lexicographic order is training order; the step is unpadded, and score === nothing omits the score segment. The score is printed with %.6g and needs no sanitizing, since check_control_readback made it a finite Float64; the metric name goes through ReactantNitro.sanitize_metric_name. A method must accept kwargs... so a keyword added later is not a MethodError:

ReactantNitro.checkpoint_filename(e::MyExp; epoch, score, kwargs...) =    "e$(lpad(epoch, 4, '0'))_$(round(something(score, 0.0); digits = 4)).jld2"

Checkpointer's name, bound at setup, calls this on every write, so a redefined method takes effect at the next checkpoint of a live run.

checkpoint_info(path) -> NamedTuple

Read a checkpoint's metadata without its weights. The fields are format_version, framework_version, epoch, step, seed, preset, the logger's run_id, run_url and logger_type, stop_reason, anchor_checksum, and metrics, the validation metrics of the epoch it was written at.

i = checkpoint_info("runs/MyExp/epoch-0012-step-116520-mae=0.0253034.jld2")i.epoch, i.step, i.preset, i.metrics.mae, i.run_id

Use this instead of opening the JLD2 file by hand: a record is one object whose four weight-bearing fields load with it, and printing it carelessly prints the whole parameter tree. Opened in a process without ReactantNitro loaded, JLD2 reconstructs a look-alike type. The checkpoint's own run id (see checkpoint_run) is in the manifest, not here. For a whole directory, use read_manifest instead of calling this in a loop.

checkpoint_run(nitro) -> String or nothing

The id this handle's checkpoints are written under, the run of their manifest entries, and what names the run in a checkpoint source: weights = "<id>" => :best. nothing with no Checkpointer, and for a run continuing the unnamed run of a manifest written before runs were recorded, which runs lists as "legacy". Not run_id, which is the logger's.

checkpointer(e)

The run's checkpointer. Default _field(e, :checkpointer, Checkpointer()); also a Nitro keyword. nothing disables checkpointing. Called once per construction, and not at all when the keyword is given.

Checkpoints rank by the run's monitor; set that, not a metric here:

ReactantNitro.monitor(::MyExp) = :f1 => :max

compile_view(e) -> e_trace

The view every trace site sees: each Host field replaced with a StrippedHost sentinel, everything else passed through, handed to the tracer as Const. The real e is used everywhere outside the trace. @experiment generates a type-stable method per experiment; the fallback here covers a hand-written one. It leaves Device fields in place, which is right for every trace the framework runs; the exported program uses export_view instead.

config_metadata(::Type{E}) -> NamedTuple

One entry per field of E, in declaration order:

(; aux_weight = (; type = Float32, default = 0.25f0, kind = :device,                   doc = "Weight of the auxiliary ...", line = 3), ...)

kind is :device, :host or :graphconst; type is the declared type; doc and line come from the declaration; default is ReactantNitro.NO_DEFAULT for a required field. The fallback synthesizes a record for a hand-written experiment from fieldnames and fieldtype, with doc = nothing, line = 0 and no defaults.

current_epoch(nitro) -> Int

The current epoch, 0 before the first one begins.

current_step(nitro) -> Int

The number of optimizer steps taken. With accum = N each optimizer step consumes N micro-batches, so this count is a factor of N below the micro-batch count. Schedules are indexed by it.

decay_anchor(e, ::Val{group}) -> :zero | :w0 | AbstractArray

What this group's decay pulls toward. Default :zero, ordinary weight decay. :w0 decays toward the parameters as build_model returned them, which is L2-SP (Li, Grandvalet and Davoine, ICML 2018, arXiv:1802.01483). An explicit array, aligned with the group's flat parameter slice, anchors to values other than the initialization. :w0 on a randomly initialized group anchors it to a random point, and the framework cannot tell which groups are pretrained, so :zero is the default. Anchoring puts a per-group checksum in the checkpoint record, which resume verifies; the anchor arrays are not stored, so an anchored run whose init is not deterministic cannot be resumed.

ReactantNitro.default_no_decay(keypath, param) -> Bool

The built-in decay exclusion: every 1-D parameter, which in Lux is biases and every normalization layer's scale and shift. Exported so no_decay can extend it: ReactantNitro.no_decay(::MyExp, ks, x) = default_no_decay(ks, x) || (:fc in ks).

derive(e, data) -> NamedTuple

Field values computed from the dataset, merged into the experiment. Default (;). Every key must name an existing field. Runs before device conversion, so it sees the pre-conversion experiment and returns host values; a derived Device field stays out of the cache key and a derived GraphConst field is baked into it. Derive numerics and keep structural values as constants, since a derived shape makes the compiled program depend on the data. Resume recomputes derived values; weights from a checkpoint record restores them and skips derive.

derive(e::MyExp, data) = (; class_weights = inv_freq(data))

device_fields(::Type{E}) -> NTuple{N,Symbol}

The names of E's Device fields, in declaration order. Generated by @experiment; defaults to () for a type that declares none.

Defining this by hand, alongside host_fields and config_metadata, is the documented escape hatch for an experiment written without the macro.

device_value(nitro, name) -> value

The host value of a Device field of the live experiment, which holds device values after setup step 5. The counterpart to set_device!. Errors on a field that is not Device, since a Host or GraphConst field is already a host value.

dispatch_variant(e, field) -> Val

Lift a Symbol config selector to a type, so variant branches below the barrier fold at trace time. Default Val(getproperty(e, field)).

early_stop(e)

The run's early-stopping rule. Default _field(e, :early_stop, nothing), no early stopping; also a Nitro keyword. A rule is checked once per epoch after validation, so setup refuses one on a handle with no val split.

ema(e) -> EMA or nothing

The run's exponential moving average of the weights. Default _field(e, :ema, nothing), no EMA; also a Nitro keyword. See EMA.

eval_loss(e) -> Bool

Whether validation and evaluation also report the loss when the experiment defines its own metrics. Default _field(e, :eval_loss, false); also a Nitro keyword. When on, each eval batch also runs loss in its own small compiled program, on the device as in training, and loss and each loss/<term> it returns are added to what finalize_metrics returns, averaged over batches. A split whose batches lack the fields loss reads is skipped. A key that metrics also returns is an error naming both sources. Without metrics the framework reports the loss anyway, so the flag changes nothing there.

Prefer it to calling loss inside metrics: under the default :host residency that runs loss on the CPU, outside the program that trains, with the experiment's Device fields still on the device.

evaluate(nitro; split = :test) -> NamedTuple

Run any named split from the data collection, erroring on a name build_data did not return and listing the available ones. Shares one compiled forward and one compiled metrics with validation and inference, so moving between them never recompiles. Runs on a worker thread; Ctrl+C stops at the next batch boundary.

experiment(nitro) -> e

The experiment after device conversion, which every hook except build_data and derive sees. Its Device fields hold device values.

export_client_inputs(e) -> Union{Vector{ExportSpec},Nothing}

Optional. Client-facing input specs, used to put axis_letters (e.g. "whcn") in the manifest, which the traced specs cannot carry. Each spec needs a dtype and a shape. With no method (the default, nothing) the server uses the executable input specs, which are the wire inputs. Requires export_postprocess, as export_client_outputs does.

export_client_inputs(e::Refiner) = [    ExportSpec("img", UInt8, [e.w, e.h, 1, 1]; axis_letters = ['w', 'h', 'c']),    ExportSpec("g0", Float32, [e.dims, e.k, 1]; axis_letters = ['g', 'k']),]

export_client_outputs(e) -> Union{Vector{ExportSpec},Nothing}

Optional. What a client receives after export_postprocess has run, when that differs from what the executable emits. Refused without a postprocess, since the server would reject the bundle. The framework cannot derive these, so dtype and shape are required. batch_axis is not checked, though batch-last is still the convention, and axis_letters and -1 axes are allowed. The default is nothing.

export_client_outputs(e::CropClassifier) = [    ExportSpec("region_prob", Float32, [e.num_classes, 1]; batch_axis = 2),    ExportSpec("region_logits", Float32, [e.num_classes, 1]; batch_axis = 2),]

export_inputs(e) -> Vector{ExportSpec}

Required. The tensors a client sends, in the order it sends them. The order fixes the positional order the program is traced with and the tensor names in the bundle, so reordering the specs changes the wire contract. Each spec needs a dtype and a shape, from which the framework builds the example arrays it traces, and the batch axis must be last; its size is a placeholder replaced by each of export_model's batch_sizes.

export_inputs(e::CropClassifier) = [ExportSpec("img", UInt8, [e.sz, e.sz, 1, 1])]

export_model(nitro, backend; dir, name, batch_sizes = [1],
            provenance_root = nothing, provenance = Dict()) -> String

Export a trained model to backend and return the path written. It takes a Nitro rather than a checkpoint path, since the handle already carries the restored weights, the experiment, the preset and the seed.

using ReactantServerExport                       # the extension that provides the backendnitro = Nitro(e; weights = :latest => :best, data = (;))export_model(nitro, ReactantServerBundle(); dir = "export_out", name = "my_model_v1")

data = (;) skips build_data, which export does not need: it reads weights and traces a graph. Each entry of batch_sizes is a separately compiled program. Export is a CPU trace and requires a single-device handle.

Provenance is merged from five layers, lowest precedence first:

Without provenance_root the export still succeeds, but the bundle omits the commit, tree hash and working-tree patch, so it cannot say which code produced it.

In order: check the hooks against each other and the batch-last rule, build example wire arrays, check that nothing traced is device-resident, run the program eagerly once to verify the output shapes, merge the provenance, and call write_export under the ExportCompiling phase. Like every entry point it runs off the interactive thread when it can; ^C stops the wait, and the export finishes in the background rather than leave a partial bundle.

export_outputs(e) -> Vector{ExportSpec}

Required. Which leaves of forward's output ship, and under what names. forward often returns more than a client needs, and selecting a subset here needs no second program. Each spec's source (its from, else its name) must be a key of the NamedTuple forward returns or a wire input; when forward returns a bare array, there must be exactly one spec. dtype and shape are derived from the trace; if given, they are checked against it, excluding the batch axis.

from lets a bundle's tensor name differ from its source, so export needs no change to forward and no recompile of the training program. It covers a rename (state_out from forward's g_pred) and an echo: a wire input returned as an output, because a serve-time postprocess receives only outputs.

export_outputs(e::Refiner) = [    ExportSpec("state_out"; from = :g_pred),   # rename: a leaf `forward` returns    ExportSpec("path_len"),                    # the usual case: the name is the source    ExportSpec("valid"; from = :valid),        # echo: a wire input, returned to the client]

An echo must declare from: a spec whose source is a wire input is refused without one. The echoed value is the wire value, not the preprocessed one.

export_postprocess(e) -> Union{String,Nothing}

Optional. The source of the model.jl that ships in the bundle, for work that does not belong in a compiled graph, such as a softmax or a decode. The framework never runs it. The default is nothing, which ships no model.jl. When the postprocess changes what a client receives, declare that with export_client_outputs or export_client_inputs; both are refused without a postprocess.

export_preprocess(e, wire...) -> NamedTuple

Optional and traced. Converts the wire inputs into the batch forward takes, inside the compiled program. It receives one positional argument per export_inputs spec and returns the NamedTuple batch forward declares its keywords against. With no method, each wire input becomes the batch field of the same name. It follows the same rules as any traced hook.

It is also called eagerly on host arrays, by export_model's routing and verification probe and by the backend's shape discovery, so it must work in both modes. The usual conversion does not: Float32.(img) has no traced broadcast method for an integer array, so it needs two methods:

_wire_to_f32(x::AbstractArray) = Float32.(x) ./ 255.0f0_wire_to_f32(x::Reactant.TracedRArray) =    Reactant.Ops.convert(Reactant.TracedRArray{Float32, ndims(x)}, x) ./ 255.0f0export_preprocess(e::CropClassifier, img) = (; img = _wire_to_f32(img))

export_provenance(nitro) -> Dict{String,Any}

The framework's provenance entries for the handle's weights: framework, reactantnitro_version, seed, run_dir, the flat config, and preset when there is one. Repository state is site_provenance's job.

checkpoint names the file the restore read (Nitro.checkpoint_source) and is omitted for fresh weights. averaged_from is present for averaged weights and names what was averaged: the spec, run, directory, files, epochs and weights. weights_from describes a warm start from another handle (weights = other::Nitro): its experiment, epoch, step and run directory. trained_run_id and trained_run_url come from the restored record, not this handle's logger, which on a weights = path construction names the exporting process; each is omitted when unknown. The checkpoint's epoch and metric are not included.

export_provenance_extra(e; checkpoint = nothing) -> Dict{String,Any}

The model's own provenance entries, merged over the framework's and the site's; the default is empty. Use it for facts a consumer of this model needs and no tensor shape shows: the labeling variant, the class names in order, the crop geometry. It is a hook so that every export includes it without the caller having to pass it.

checkpoint is the file the weights came from (Nitro.checkpoint_source), or nothing for fresh weights. With no checkpoint, a model should report any claim about its training data as unverifiable.

function ReactantNitro.export_provenance_extra(e::MyExp; checkpoint = nothing)    prov = Dict{String, Any}(        "model" => "MyModel",        "class_names" => collect(String, class_names(e.variant)),        "crop_sz" => [e.sz, e.sz],    )    checkpoint === nothing || (prov["trained_from"] = abspath(String(checkpoint)))    return provend

These keys override the framework's and the site's; export_model's provenance argument overrides them.

export_view(e) -> e_export

The experiment as an exported trace sees it: compile_view(e), which drops every Host field, with every remaining device-resident value copied to the host.

In training, a Device field is a traced input. An exported bundle is called as executable(inputs..., weights...), and Reactant lifts every device-resident value a traced closure reaches into an extra argument the bundle cannot supply, so the bundle would fail every inference with an argument-count mismatch. Copying to the host avoids that: a Device field forward reads is baked into the graph at its current value, and one it does not read is unused. The view cannot drop the unread fields instead, since nothing declares which fields forward reads. A large baked-in array is stored once per batch size; an experiment holding one should override this method.

finalize_metrics(e, acc, split) -> NamedTuple

Turn one split's combined metrics into the reported values. Runs on the host, once per split per evaluation; the default returns acc unchanged. acc has the shape of metrics' return with each value combined across batches as its mode says, and split is the split's name, such as :val. Use it for a metric that is not a mean of per-batch values, and to branch on split, which costs no compile here:

finalize_metrics(e, acc, split) = (; f1 = 2acc.tp / (2acc.tp + acc.fp + acc.fn))

finish!(lgr, status) -> nothing

Close the run on the backend. status is :completed, :early_stop, :requested (a request_stop!), or :error, matching the stop_reason recorded in the checkpoint.

Do not map Base.close to finish!(:completed): in a finally it runs after the driver's finish!(:error) and would mark every failed run completed.

forward(e, model, ps, st; <declared batch fields>) -> (outputs, st_new)

Required. The model's computation, from inputs to outputs. Training traces forward then loss, validation traces forward then metrics, and predict runs forward alone, which is also the exported program. Declare inputs only, never targets, so predict can call it without them.

Declare exactly the batch fields you want as keywords; the framework resolves the method once at setup and passes that subset. A keyword with a default is an optional batch field. A method ending in kwargs... receives the whole batch unchecked.

loss, metrics and train_metrics receive outputs, the first element of the return, as one positional argument. For more than one output, return a NamedTuple, which is also what export_outputs names leaves by:

forward(e, model, ps, st; tokens) = ((; logits, energy), st_new)loss(e, out; target) = ce(out.logits, target) + e.w * mean(abs2, out.energy)

Every leaf of the output tree must be an array with the batch dimension last: the short final eval batch is padded and sliced back along the last axis, and a scalar already reduced over the batch has nothing to slice. Always return the state (st itself for a stateless model). The framework switches between train and eval mode and discards the eval-mode st_new.

from_preset(E::Type, name::Symbol; overrides...) -> e

Build an experiment from presets(E)[name], with overrides winning over the preset. Every preset key is validated against fieldnames(E). A preset is field values, so the marker semantics are untouched: switching presets recompiles exactly when a GraphConst field differs. This does not record the name; use Nitro(E, name; ...) for that.

e = from_preset(MyExp, :current; max_epochs = 40)

grad_rule(e) -> GradRule
grad_rule(e, hp) -> GradRule

Which gradient each micro-batch contributes: SingleGrad by default, or a PerturbRule such as SAM. Its fields are host numbers, never Device fields.

The one-argument form returns the rule, and a scheduled field (rho = ... or grad = (; rho = ...)) is filled in by name. The two-argument form mirrors a Level 2 optimizer: hp holds this step's scheduled rule fields by name and bind values as hp.bind.<name>, all host numbers, and you place them:

ReactantNitro.grad_rule(::MyExp) = SAM(0.05f0)ReactantNitro.grad_rule(::MyExp, hp) = SAM(0.05f0 * hp.bind.sam_warm)

Called on the host once per optimizer step; the framework moves the rule's fields to the device, so a changing value never recompiles. The rule's type must not change. Manual mode supports SingleGrad only.

gradient_clip_norm(e) -> Real

The global-norm gradient clip threshold. Default _field(e, :gradient_clip_norm, 0f0), where 0 means off; also a Nitro keyword. Applied to the fully accumulated gradient at the start of the optimizer program, not as a chain member, which would clip per group. It is a trace-time host constant, so a disabled clip emits no ops and changing it recompiles the optimizer program only. A traced threshold could not make 0 mean off, since 0/0 on an all-zero gradient yields NaN. The name matches Lightning's.

history(nitro) -> MetricHistory

The per-epoch series this handle's train! calls produced, as a MetricHistory. Each handle keeps its own: a resumed handle's history begins at the first epoch it trained, and its table notes that earlier epochs ran in another process. Nothing is read from disk, so it works with any logger backend or none.

host_fields(::Type{E}) -> NTuple{M,Symbol}

The names of E's Host fields, in declaration order. Generated by @experiment; defaults to () for a type that declares none. This is what compile_view reads.

An unmarked field is Host, so it is listed here. The GraphConst fields are the complement, setdiff(fieldnames, device_fields, host_fields).

lambda(e) -> Real
lambda(e, ::Val{group}) -> Real

The decoupled weight decay coefficient, per group. Default 0, which adds no Decay to the chain. Each group decays toward its own decay_anchor, so :w0 on a pretrained backbone and :zero on a fresh head share one rule and one opt.lambda schedule. The value reaching the rule is multiplied by the group's effective learning rate. Scheduling opt.lambda while also defining per-group lambda accessors is a setup error.

learning_rate(e) -> Real
learning_rate(e, ::Val{group}) -> Real

The base learning rate (default 1f-3) and the per-group bases (default learning_rate(e)). Per-group accessors define ratios and the schedule sets the default group's absolute value: η_g(t) = eta_sched(t) * (learning_rate(e, Val(g)) / learning_rate(e)). learning_rate(e) == 0 is a setup error, since the ratio is undefined.

load_checkpoint(ckpt, path) -> CheckpointRecord

Load the record at path through the checkpointer that wrote it. It dispatches on the checkpointer rather than the path so that a checkpointer with a non-file backend can implement it. The Checkpointer method migrates records written before a field was added and refuses a record of another format_version. With checkpointer = nothing it returns nothing.

log_confusion!(lgr, matrix, labels; epoch) -> nothing

Log a confusion matrix, given as a plain Matrix{Int} that the backend adapts. The framework never calls this, since only the experiment knows whether it has a confusion matrix and what its labels mean. Accumulate it in metrics with :sum, stash it from finalize_metrics, and emit it from a phase monitor on the transition out of EvalStepping:

metrics(e, out; y)            = (; confusion = confusion_matrix(out, y) => :sum)finalize_metrics(e, acc, spl) = (e.rt.last_confusion = acc.confusion; derived_scalars(acc))# then, in a phase monitor:log_confusion!(info.logger, e.rt.last_confusion, e.rt.class_names; epoch)

log_every_n_steps(e) -> Int

Write the train line only on optimizer steps that are a multiple of n. Default _field(e, :log_every_n_steps, 1); also a Nitro keyword. The train line carries the loss, train_metrics and the schedule values; thinning it keeps a fast run from flooding its logger. Only the writing is thinned: validation and the per-epoch lines are unaffected, the loss is still read back every micro-batch for the divergence check and history, and train_metrics still runs every micro-batch.

log_metrics!(lgr, metrics; step, epoch, context, kwargs...) -> nothing

Log a NamedTuple of metrics. context is one of:

• "train": one line per logged optimizer step (every log_every_n_stepsth step), carrying the loss, the train_metrics values and, under autolog_schedules, the schedule values under schedule/ keys
• "validate": one line per epoch with the val split's metrics
• "data": one line per epoch with data_wait_frac, the fraction of the epoch the host spent waiting on data; a separate context so it does not mix into the other two series

Values arrive as host values. The framework drops non-finite Real values before calling and passes other values (arrays, tuples) through unchanged. A method must accept kwargs... so a keyword added later does not break it.

log_other!(lgr, key, value) -> nothing

Log one free-form value under key. Setup logs the schedule binding report this way, as log_other!(lgr, "binding_report", str).

log_params!(lgr, params) -> nothing

Log the run's hyperparameters. Setup calls this before the first compile, so a crash during compilation still leaves them logged. params is a flat NamedTuple of the experiment's configuration and the run's settings (seed, accum, max_epochs, the resolved prefetch settings), plus preset when the run names one. Keep keys flat: dotted keys split comparison tables in hosted trackers.

log_tags!(lgr, tags) -> nothing

Free-form tags for the run; the backend decides what a tag is.

logger(e)

The run's logger. Default _field(e, :logger, JSONLogger()); also a Nitro keyword. nothing turns logging off. Called once per construction, as the keyword's default, so at the start of construction before build_data, and not at all when the logger keyword is given. A side effect in the accessor (opening a file, registering a run with a hosted tracker) therefore happens once per Nitro; a field default instead runs when the experiment is constructed. The default JSONLogger opens its file lazily, after setup has pointed it at the resolved run_dir.

logger_info(lgr) -> NamedTuple

The parameters that identify where the backend writes, as a plain NamedTuple: an experiment key and URL, a run id and project, or for JSONLogger its file path. Informational, like run_id; the nitro_logger tool renders it. Default (;). Return plain data, never the live object, since the value is passed to tools and reports. Not stored in the checkpoint record.

logger_info(lgr::JSONLogger) -> NamedTuple

(; path), the file the logger writes to, or nothing before setup has given a pathless logger its run_dir.

logger_info(nitro) -> NamedTuple

The identifying fields of the run's logger, logger_info(nitro.logger), readable while the run trains. The default JSONLogger reports its path; a hosted backend reports fields such as its URL, experiment key and workspace. The Kaimon nitro_logger tool renders this.

logger_state(lgr) -> Any

The state a resumed run needs to continue the same logged run, such as a hosted experiment key, a run id plus project, or a tracking URI. Opaque to the framework and stored in the checkpoint record. Default nothing: no resumable state, so nothing is stored and no reattachment is attempted. Return plain serializable data (String, NamedTuple, Dict), never the live backend object, since the record is written with JLD2 and outlives the process. A logger that returns anything else must define reattach!. run_id is separate and informational only.

loss(e, outputs; <declared batch fields>) -> scalar or (scalar, terms::NamedTuple)

Required. The scalar Enzyme differentiates, traced together with forward. outputs is already the first element of forward's return; do not unpack it again, since first(outputs) is the first element of the prediction array and the run would train on one sample without an error. When an experiment defines no metrics, the default validation metric calls loss with the same declared batch fields.

Returning (l, terms) also logs each term, a flat NamedTuple of scalars, as loss/<term>: on the train line from the gradient program itself, so nothing is recomputed, and wherever the validation loss is reported (see eval_loss). Only l is differentiated.

ReactantNitro.loss(e::MyExp, ŷ; y) = (ce = crossent(ŷ, y); kl = kldiv(ŷ); (ce + 0.1f0 * kl, (; ce, kl)))

manual_training(e) -> Bool

The manual-mode flag. Default: whether the experiment type has a train_step method. Override to false to keep the method in source while using the automatic loop. Read once at construction and frozen; flipping it on an existing handle is reported by fixed_config_report.

max_epochs(e) -> Int

How many epochs to train for. Default _field(e, :max_epochs, 1), so a bare experiment trains one epoch. Also a Nitro keyword, the usual way to raise it on resume. Resolution order: the keyword, a user method, a field on e, the default. It runs on the host against the real e, so a max_epochs field must not be GraphConst; leave it unmarked (Host). Setup checks this.

metrics(e, outputs; <declared batch fields>) -> NamedTuple of `value => mode`

The validation metrics of one eval batch. Each metric is value => mode, where the mode says how batches combine:

metrics(e, outputs; lab) = (; err = sum_abs_err => n_items,                              cm = confusion => :sum,                              ranking = (; score = p => :concat, n_pos = count(lab) => :sum))

A key may be a group (NamedTuple or Tuple). A mode on a group covers every leaf below it, and each leaf of a group must be covered by exactly one mode. finalize_metrics receives the same shape with the modes removed. A :concat leaf must have one entry per sample on its last axis, and must not be returned from finalize_metrics unreduced. The tuple (value, mode) is the older spelling of value => mode and still works at the top of a key.

metrics never sees padding. When the experiment defines none, the framework reports loss = loss(e, outputs; ...) => 1, the mean over batches, on every split. Name metrics without a split (acc, not val_acc): the split is the context they are logged under. metrics takes no split argument, since that would compile one program per split; branch on the split in finalize_metrics.

metrics_residency(e, hook::Symbol) -> :host | :device

Where a metric hook runs. hook is :metrics or :train_metrics; the defaults are :host for metrics and :device for train_metrics. train_metrics runs once per micro-batch, where transferring the full outputs would cost far more than the scalars it returns. metrics runs once per eval batch, where the transfer is cheap, and on the host it can use any Julia code (matching, sorting, a library that knows nothing about Reactant).

ReactantNitro.metrics_residency(::MyExp, ::Symbol) = :device    # a traceable validation metric

For the hook author, the difference is whether outputs holds device or host arrays. A :host hook is part of no program, so editing it recompiles nothing.

monitor(e) -> Pair{Symbol, Symbol} or nothing

The validation metric the run is judged by, as metric => mode with mode :min or :max. Default _field(e, :monitor, nothing); also a Nitro keyword. Left at nothing, setup uses :loss => :min when the validation loss is guaranteed to exist: the experiment defines no metrics, or eval_loss is on. Otherwise there is no default, and setup refuses a checkpointer with a Best selector, an EarlyStopping or a Plateau that would follow it. The Checkpointer, EarlyStopping and every Plateau watch it unless given their own metric and mode, so it is set in one place:

ReactantNitro.monitor(::MyExp) = :acc => :max

metric is a key of the validation metrics: what finalize_metrics returns (by default, what metrics returns), or :loss, which the framework supplies when the experiment defines no metrics. Metric names carry no split, since the split is the context: the validation loss is :loss and an accuracy is :acc, never :val_acc, and a logging backend's display name (validate_acc) is not the key. With a custom metrics the key cannot be checked before validation runs, so a misspelled one raises at the end of the first epoch. :val_loss, the name before 0.2.0, is accepted as :loss with a warning until 1.0.0.

n_devs(e) -> Int

The number of local devices to shard the batch over; also a Nitro keyword. 1 skips the mesh. Default _field(e, :n_devs, length(Reactant.devices())), every visible device, so set CUDA_VISIBLE_DEVICES to restrict a run; with no GPU visible it is 1. A session pin from setup_devices! (the nitro_setup tool) takes precedence over the field and the default, and an explicit n_devs keyword on Nitro takes precedence over the pin. The batch size is global and is split across the devices, so it must divide by n_devs.

no_decay(e, keypath::Tuple, param) -> Bool

Whether this parameter leaf is excluded from weight decay. Default default_no_decay. It receives the leaf as well as its keypath, so it can test ndims and size; it runs on the host, once at setup.

ReactantNitro.no_decay(::MyExp, ks, x) = default_no_decay(ks, x) || (:fc in ks)   # extendReactantNitro.no_decay(::MyExp, ks, x) = (:fc in ks) && (:weight in ks)          # replaceReactantNitro.no_decay(::MyExp, ks, x) = false                                   # decay everything

An excluded leaf gets a mask of 0, and Decay computes mask * lambda * (x - anchor), so exclusion overrides the anchor. Exclusion is per leaf while the coefficient and anchor are per group, because biases occur in every group and splitting groups for them would also split the learning-rate ratio. The mask is a device buffer passed as a value, so it is not in the compile cache key and changing it recompiles nothing. Resolved once at construction.

no_perturb(e, keypath::Tuple, param) -> Bool

Whether a PerturbRule such as SAM leaves this parameter leaf unperturbed and out of ‖g‖. Default false: every leaf is perturbed, as in Foret et al. Independent of no_decay; to mirror it, no_perturb(e::MyExp, ks, x) = no_decay(e, ks, x). Runs on the host once, at construction.

nonschedulable(::Type{R}) -> NTuple{N,Symbol}

The fields of rule type R that may not be scheduled. to_device_rule also leaves them on the host, so one declaration decides both. Default (). The shipped methods:

nonschedulable(::Type{<:Optimisers.Adam})     = (:beta, :epsilon)nonschedulable(::Type{<:Optimisers.RAdam})    = (:beta, :epsilon)nonschedulable(::Type{<:Optimisers.AdamW})    = (:beta, :epsilon, :couple)nonschedulable(::Type{<:Optimisers.ClipNorm}) = (:p, :throw)nonschedulable(::Type{<:Decay})               = (:anchor, :no_decay_mask)

beta is a Tuple; p is structural and promoting it breaks _norm's ::Real dispatch; any parameter-sized field is non-schedulable. Declare one method per concrete rule type, never a Union: a more specific method shadows a Union method with no ambiguity warning.

optimizer(e) -> Type{<:Optimisers.AbstractRule}
optimizer(e, ::Val{group}) -> Type{<:Optimisers.AbstractRule}
optimizer(e, group::Symbol, hp) -> Optimisers.AbstractRule

The optimizer rule, at one of three levels of control.

Level 0: declare nothing and get RAdam, with per-group hyperparameters from learning_rate and lambda. RAdam rectifies the adaptive variance term, so no hand-tuned warmup is needed; it needs Optimisers 0.4.8, the first release whose Reactant extension traces it.

Level 1: return a rule type, for every group or per group, and the framework constructs it, passing the resolved values by field name. The Decay tail, the per-group accessors and the no-decay mask still apply. A rule with its own lambda field (AdamW) is rejected here, since it would decay twice.

optimizer(::MyExp) = Optimisers.Momentumoptimizer(::MyExp, ::Val{:backbone}) = Optimisers.Momentum

Level 2: build the chain yourself from the group and its resolved, device-converted hp. It is called at setup and again every optimizer step:

optimizer(::MyExp, group::Symbol, hp) = Optimisers.OptimiserChain(    Optimisers.ClipNorm(1.0f0, 2.0f0, false),    Optimisers.Momentum(hp.eta, get(hp, :rho, 0.9f0)),)

hp carries eta (the group's effective learning rate), lambda (already multiplied by it, and nothing when it is zero and not scheduled), anchor, no_decay_mask, and one key per scheduled rule field, present from the first call. A scheduled field reaches the rule only if the factory reads it from hp. Setup checks this on the host: a scheduled field whose value does not change the returned chain is a setup error, and the binding report shows where each scheduled field landed. No level accepts a prebuilt chain with fixed values, since the rules carry device scalars that are rebuilt every step.

param_group(e, keypath) -> Symbol

The parameter group a leaf belongs to, as a function of its keypath. Default :default for every keypath, a single group. Each group gets its own optimizer chain, learning-rate ratio, decay and anchor. :default is group 1 and the rest are numbered in order of first appearance, which bakes into the program.

param_group(::MyExp, ks) = ks[1] === :backbone ? :backbone : :default

parameters(nitro) -> ps

The parameters, as the Lux tree build_model returned. After train! these are the trained values.

perturbation(rule::PerturbRule, g, p, mask) -> ε

The parameter perturbation for one chunk. g (the unperturbed gradient), p (the parameters), mask (1 for a leaf that may be perturbed, 0 for one no_perturb excludes) and ε are NTuple{G}s of flat per-group buffers. Norms must reduce over the parameter axes only, so a per-chunk variant can map the method over chunks unchanged. Traced.

perturbation(rule::ASAM, g, p, mask, adaptive) -> ε

ASAM's perturbation. adaptive holds 1 where adaptive_perturb scales the leaf by |w| + eta and 0 where it is left unscaled.

phase(nitro) -> Phase

The run's current phase: the latest transition, which monitors also receive.

predict(nitro, batch::NamedTuple) -> outputs
predict(nitro, loader) -> iterator

Inference: forward alone, in eval mode.

Takes a batch NamedTuple or anything iterating them, never a bare array: forward is keyword-routed from the batch's field names, and the error for a bare array names the fields it declares. Only forward's fields are required, so a prediction batch needs no labels.

One batch in, one output tree out, sliced to the real sample count, as host arrays. A loader in, a lazy iterator out, one element per batch; collect it for all of them. An output leaf whose last dimension is not the batch raises rather than slicing the wrong axis.

presets(::Type{E}) -> NamedTuple

The table of named configurations for an experiment type. Default empty. An entry sets some fields and leaves the rest at the struct's defaults; to build one preset on another, merge the NamedTuples.

ReactantNitro.presets(::Type{MyExp}) = (    reference_v1 = (; n_classes = 8, batch_size = 16, rotate_deg = 7.5),    current      = (; n_classes = 10, batch_size = 32, rotate_deg = 12.0),)

Presets let the checkpoint record and the logger name the recipe behind a result, and make a typo in a key an error instead of a no-op. They are not a reproducibility mechanism; provenance is.

ReactantNitro.progress_counter() -> Int

A monotonic count of units of work this process has completed: one per optimizer step in training, one per batch in evaluation. Only changes in it are meaningful. A run stays in TrainStepping for a whole epoch, so a stall watchdog compares this count across two observations to measure time since the last completed unit. It counts evaluation batches because nothing on the handle advances during an eval loop.

progress_note!(msg) -> nothing

Show msg beside the current phase of the progress display, as setup [building data] (decoding 1200/5000). Inside with_progress_note it replaces that block's note; outside one it sets a note that lasts until the next phase or stretch. Cheap enough to call per item: it redraws at most ten times a second.

rank(dist) -> Int

This process's rank. dist is always nothing today and rank(nothing) == 0.

read_manifest(dir) -> Vector

Which checkpoints a run directory holds, without opening one. Each entry carries file (a basename), epoch, score (the retention metric, or nothing for a run with no val split), stop_reason (nothing until the run records how it ended), run (the id of the run that wrote it), metric and mode (what score measures, and which way is better), written (the wall-clock time() of the write) and parent (for a run that branched from an earlier record, that record's path). A manifest written before runs were recorded reads with run, metric, mode and parent as nothing and written as 0.0: one unnamed run, older than every named one. Returns an empty vector for a directory no run has written to. Use checkpoint_info only for what the manifest does not carry.

for e in sort(read_manifest("runs/MyExp"); by = e -> (e.written, e.epoch))    println(e.run, "  ", e.epoch, "  ", e.score, "  ", e.file)end

reattach!(lgr, state) -> nothing

Point a logger back at the logged run that state (from logger_state) identifies. Setup calls it on resume, before any metric is logged. There is no default method; a logger needs one exactly when its logger_state can return something other than nothing. Resuming into a logger of a different type than the one that wrote the record is refused, naming both.

register_phase_monitor!(f) -> handle
register_phase_monitor!(nitro, f) -> handle

Register f to be called on every phase transition. The module-level form is copied into each handle at its next entry point, including handles already built; the nitro form adds f to that handle's own registry and takes effect immediately. Heartbeats and watchdogs are written against this; the framework ships neither.

f is called as f(phase, step, epoch, info) with info::NamedTuple, since do-block functions cannot take keyword arguments:

register_phase_monitor!() do phase, step, epoch, info    phase isa Compiling && @info "slow phase, do not kill me" phase    phase isa Terminal  && close(my_heartbeat)end

step and epoch are nothing until the first epoch begins. info carries at least (; nitro, logger, is_rank0), plus metrics on the transition out of EvalStepping; during Nitro(e) construction, before a handle exists, nitro and logger are nothing. An exception thrown by a monitor is caught and warned about once per monitor, and the run continues. Module-level monitors fire in registration order, then the handle's own. Unregistering a module-level monitor does not remove it from handles that have already copied it.

Monitors fire on transitions only. A compile is a blocking foreign call that emits nothing, so a monitor that must prove liveness needs its own task on the :interactive pool, and can extend its timeout while the phase is a Compiling.

render(nitro; split = :val, batches = 1, predictions = false, out_dir, tag) -> Vector{String}
render(nitro, batch; predictions = false, out_dir, tag = "") -> Vector{String}

Write one figure per sample of real batches, through visualize and save_figure, and return the paths written. The first form takes the first batches batches of split in the loader's order; the second renders the batch given.

With predictions = false, the default, visualize gets outputs = nothing and nothing is compiled or run on the device, which makes it the check of the data itself. With predictions = true it gets the outputs of predict for the batch. batches counts batches, not samples: at batch_size = 64 one batch writes 64 figures, so pass a smaller batch to the second form to write fewer. Renders of the same split compare across runs only if its loader is deterministic.

Figures go to out_dir, default joinpath(run_dir(nitro), "viz"), under the stem <tag>_<i>, where i is the sample's index zero-padded to three digits and counted across batches. tag defaults to the split name in the first form; in the second it defaults to "", giving the stem <i>.

render(Nitro(e); split = :val)                                        # the data gaterender(Nitro(e; weights = "runs/x/best.jld2"); predictions = true) # from a checkpointrender(nitro, batch; predictions = true)                              # a batch in hand

request_stop!(nitro) -> nothing

Ask the run to stop, from a REPL or from a phase monitor. The stop is graceful: training leaves the step loop at the next step boundary, then validates, writes the checkpoint, finalizes the logger and exits through Done with stop_reason = :requested. A standalone validate or evaluate stops at the next batch.

run_dir(e) -> String

The run's output directory. Default _field(e, :run_dir, joinpath("runs", string(nameof(typeof(e))))), relative to the working directory; also a Nitro keyword. Checkpoints, the manifest, the default logger's file and resume = :auto all resolve against it. run_dir(nitro) returns the directory a handle resolved to.

run_dir(nitro) -> String

The run's output directory, against which checkpoints, the manifest and resume = :auto resolve. Defaults to joinpath("runs", string(nameof(typeof(e)))).

run_id(lgr) -> Any

The backend's run identifier, for people and reports, or nothing if it has none (the default). Stored in the checkpoint record so a checkpoint traces back to its logged run. Unrelated to checkpoint_run.

run_url(lgr) -> Any

The backend's URL for the run, or nothing if it has none (the default). Informational like run_id, and stored in the checkpoint record with it.

runs(run_dir) -> Vector{RunSummary}

The runs a run directory's manifest holds, most recent first, one RunSummary each: id, how many checkpoints it keeps and their epochs, its best score and the metric and mode that score is under, when it last wrote, and the checkpoint it branched from. The entries of a manifest written before runs were recorded are one run, "legacy", listed last. An id names the run in a checkpoint source, weights = "<id>" => :best, and checkpoint_run(nitro) is a handle's own. Answered from the manifest; no record is opened. Empty for a directory no run has written to.

runs("runs/MyExp")n = Nitro(MyExp(); run_dir = "runs/MyExp", weights = first(runs("runs/MyExp")).id => :best)

save_checkpoint!(ckpt, epoch, metrics, snapshot) -> Bool or nothing

Write a checkpoint of epoch, if the checkpointer decides the epoch warrants one. metrics is the epoch's validation metrics, empty for a run with no val split, and snapshot holds the record's other fields as host values. Returns false when the checkpointer declined; any other value, nothing included, means it wrote. train! uses this to decide whether the final stop_reason rewrite has a record to rewrite.

Checkpointer returns true or false, declining an epoch no selector in keep admits (with a Last, every epoch is admitted). It writes through with_io_retry to a temporary path renamed into place, then updates the manifest, and only then deletes the files rotation displaced. The experiment is not an argument: the filename comes from the checkpointer's name, so the object holding device values never reaches the code that serializes.

save_figure(e, fig, stem) -> path

Write what visualize returned and return the path written, as a string. e is in the signature so a model package's method dispatches on a type it owns instead of pirating Makie.Figure. stem is a path without an extension; the method chooses the format and appends it. render returns these paths, so a method that writes a video or a directory of frames returns that path. There is no default method.

function ReactantNitro.save_figure(::MyExp, fig::Makie.Figure, stem::AbstractString)    path = stem * ".png"    Makie.save(path, fig; px_per_unit = 2)    return pathend

schedules(e) -> NamedTuple

The values that vary with the optimizer step, keyed by rule field or Device field. Default (;). Every entry is a factory of the horizon: setup calls f(total) once, and the returned sched(step) runs once per optimizer step. A bare Number is a constant.

schedules(e::MyExp) = (;    eta        = total -> OneCycle(total, 1f-3),               # factory of the horizon    aux_weight = _ -> (t -> max(0f0, 1f0 - t / 5000)),         # ignores the horizon    lambda     = 1f-4,                                         # a bare Number is a constant)

Also a Nitro keyword, which replaces this table whole; to extend it, pass merge(schedules(e), (; ...)). Unlike the other run accessors, it reads no field of e. A key is a Device field name or a rule field name (eta, not lr); device and opt are reserved at the top level to qualify a name that is both:

schedules = (; eta = total -> OneCycle(total, 1f-3),               device = (; lambda = _ -> t -> 0.5f0),    # e.lambda               opt     = (; lambda = _ -> t -> 1f-4))     # the decay coefficient

A nested opt key is a parameter-group path: opt = (; backbone = (; eta = ...)) binds the :backbone group's chain only, with that group's learning-rate ratio still applied. Paths are one level, since groups are flat. In manual mode a path is a path into opt_state and its values are absolute. step is the optimizer step, and total = max_epochs * div(length(train), accum), which divides exactly because setup requires it. Schedules run on the host, so ordinary control flow is fine; one that closes over training state does not resume exactly. ParameterSchedulers.jl works here and is not a dependency.

seed(e) -> Int

The global RNG seed, set before build_data. Default _field(e, :seed, 42); also a Nitro keyword. A seed field must not be GraphConst, so that a seed sweep shares one compiled program; leave it unmarked (Host). Setup checks this. A resume restores the record's seed, with a warning if it differs.

select(sel::Selector, entries, ckpt) -> Vector{ManifestEntry}

The entries of entries that sel picks. entries is what retention decides over: this run's under scope = :run, the directory's under scope = :dir.

set_device!(nitro, name, value) -> nitro
set_device!(nitro; name = value, ...) -> nitro

Write a new value into a Device field of a live handle without recompiling, so a sweep over a device value or a changed inference threshold reuses the compiled programs.

The cache key skips device_fields when hashing and seeds on the type name rather than the parameterized type, so a new device value leaves every key component unchanged; the function asserts both before committing. Refused: a field that is not Device (a GraphConst change is a different program, and a Host field reaches no trace), a value of a different element type or size (the type would move the key; the size would not, and would fail inside XLA instead), and a scheduled field (the per-step rebuild would overwrite the write).

n = Nitro(e; weights = "runs/MyExp/best.jld2")      # weights-only, skips `derive`for thr in (0.3f0, 0.5f0, 0.7f0)    set_device!(n; threshold = thr)    out = predict(n, batch)                         # no recompileend

ReactantNitro.setup_devices!(; backend = nothing, n_devs = nothing) -> NamedTuple

Configure this process's accelerator for every subsequent run and report what is in effect, as (; backend, visible, n_devs, pinned). The Kaimon tool nitro_setup is exactly this function.

Optional: a session that never calls it runs on Reactant's default backend with n_devs equal to every visible device. backend passes through to Reactant.set_default_backend ("cpu", "gpu", "cuda", "rocm", "tpu"). n_devs pins the device count for the process, validated against the visible devices; the pin beats an experiment's declared n_devs, and an explicit n_devs keyword on Nitro still wins for that one run.

The XLA client is initialized once per process and n_devs only takes a prefix of the visible devices, so to choose which GPUs a session sees, set CUDA_VISIBLE_DEVICES before starting the process. The batch size is global and is split across the mesh, so adding devices buys throughput without changing the effective batch. Called with no arguments, this only reports.

ReactantNitro.setup_devices!()                      # reportReactantNitro.setup_devices!(backend = "cpu")       # run everything on CPUReactantNitro.setup_devices!(backend = "cuda", n_devs = 2)  # two of the visible CUDA devices

setup_optimizers(e, model, ps, st, mesh) -> opt_state

Required for a manual experiment. Build the optimizer states, typically one per network. Called once at construction, after build_model. The framework moves the result to the device and asserts that no host Number remains in any leaf's state, so an integer step counter cannot freeze under trace.

ReactantNitro.setup_optimizers(e::MyGAN, model, ps, st, mesh) = (;    gen  = Optimisers.setup(Optimisers.AdamW(1.0f-4), ps.gen),    disc = Optimisers.setup(Optimisers.Adam(2.0f-4), ps.disc))

should_stop(es, epoch, metrics) -> Bool

Update the policy with one epoch's validation metrics and return whether to stop. A custom policy implements this method for its own type; should_stop(::Nothing, ...) returns false. For EarlyStopping, an epoch improves when it beats the best so far by more than min_delta, and the run stops after patience epochs in a row without improvement. The metric is read through check_control_readback, which raises on a non-finite or non-real value.

site_provenance(backend, root) -> Dict{String,Any}

Repository state for a bundle's provenance, collected at root (export_model's provenance_root; nothing is collected without one). The default is git_provenance(root). It dispatches on the backend so a format can use its own collector: ReactantServerBundle uses ReactantServerExport.collect_provenance. A git_diff entry is a full patch, which the backend writes as a file beside the artifact rather than into its metadata.

states(nitro) -> st

The layer state, including running statistics. Plural so it is not read as the state of the whole handle.

step_optimizer(opt_state, ps, grads) -> (states, ps_new)

Run one optimizer step on a parameter tree inside a train_step closure. It calls Optimisers.update and returns the new optimizer states, without their rules, and the new parameters. The rules are removed because a replicated scalar cannot leave a compiled program on a mesh; the driver re-attaches them after the step. Call it once per network, passing that subtree:

st_g, ps_g = ReactantNitro.step_optimizer(opt_state.gen, ps.gen, g_g)

train!(nitro) -> Nitro
train!(e; kwargs...) -> Nitro

Run the training loop until max_epochs, an early stop or a requested stop, and return the Nitro. The call blocks.

train!(nitro) takes no keywords; every keyword belongs to the Nitro constructor, and train!(e; kwargs...) is sugar for train!(Nitro(e; kwargs...)). The handle is also reachable through run_ref::Ref{Nitro}, filled before the loop starts, and as info.nitro in a phase monitor. To retrain with a different stopping rule, construct another Nitro; the compile cache is module-level, so the new handle reuses every compiled program.

A non-finite loss stops the run with an error naming step and epoch; there is no rollback, and recovery is a resume from the last checkpoint at a lower learning rate. The automatic loop reads the loss back one micro-batch late, so the device runs the next micro-batch while the host waits; the error is therefore raised after at most one more micro-batch (and, at a step's last micro-batch, its optimizer step) has been dispatched, and still names the step the loss belonged to. check_divergence= false turns the stop off: the loss is logged and training continues. Every scalar the framework branches on (the loss, the checkpoint metric, the early-stopping metric) is validated on readback, because a failed BufferToHost on this stack returns garbage without raising; metrics that are only logged are not, nor is the loss under check_divergence = false.

The train line is written every log_every_n_steps optimizer steps and carries the loss, the train_metrics of the step's last micro-batch and, under autolog_schedules, each varying scalar schedule under a schedule/ key. At the end, stop_reason is :completed, :early_stop, :requested or, if the loop threw, :error.

Ctrl+C is a graceful stop: the loop runs on a worker thread, ^C becomes request_stop!, the step loop breaks at its next boundary, validation and the checkpoint still run, and the call returns with stop_reason = :requested.

A handle built with restore_best = true ends every train! that did not fail (completion, early stop, a requested stop) by loading its own run's selected checkpoint into ps and st; step, epoch, history and the last metrics stay as training ended. A further train! on that handle is refused, since those weights beside the final optimizer state and step match no state any run was in; the error lists the constructions that do. With no scored checkpoint it warns and keeps the final weights.

train_metrics(e, outputs; <declared batch fields>) -> NamedTuple of scalars

Per-step training diagnostics, computed once per micro-batch. Default (;), which compiles to nothing. At the default residency (:device, see metrics_residency) it runs inside the gradient program. Each value is a scalar logged as is, with no reduction across batches; smoothing is the logger's job. Readback is lazy, so the values cost a transfer only on logged steps.

train_step(e, model, ps, opt_state, st; <declared batch fields>)
    -> (; loss, ps, st, opt_state, stats)

Define this method for an experiment type to switch train! to manual mode; manual_training(e) = false keeps the automatic loop anyway. In manual mode you write the whole optimizer step (forward and backward passes, optimizer updates, device metrics), and the framework keeps everything around it: the batch stream and prefetch, validation, checkpointing, logging, phases, early stopping.

Called once per optimizer step, in train mode, on one device batch whose fields are routed from the declared keywords. It is traced and compiled, so it must be a named method. Compute gradients with backward(f, ps_sub, consts...); its objective must take every traced value it reads as an argument, because a traced value captured by a closure becomes a constant and yields a zero gradient with no error.

The manual-mode page works through a minimal GAN, with a generator and a discriminator on their own optimizers (see A minimal GAN).

The return is a NamedTuple, checked every step: loss (required; validated and logged), ps, st, opt_state (the new optimizer states only, since a rule's replicated device scalars cannot leave a compiled program on a mesh; the driver re-attaches the rules), and stats (device scalars logged on each train line, (;) allowed).

schedules works in manual mode, but opt keys bind into the opt_state setup_optimizers returned, since there are no parameter groups, and every value is absolute. A bare key (opt.eta) sets that field on every rule that has it; a nested key (opt.gen.eta) is a path into the tree. accum > 1 is a setup error, and train_metrics is unused, since stats replaces it. The loss hook is optional; it is needed only for the default validation loss when metrics is not defined.

unregister_phase_monitor!(handle) -> nothing

Remove a monitor by the handle register_phase_monitor! returned. The handle records which registry it was added to, so this works for both forms. Removing a monitor twice is a no-op.

validate(nitro) -> NamedTuple

Run the :val split and return the finalized metrics. The training loop runs the same pass each epoch, and validate also works on a Nitro that has never trained. Eval mode throughout; st_new is discarded, and each batch's device buffers are freed explicitly rather than left to the GC, which is the documented cause of device OOM during validation on this stack. Runs on a worker thread; Ctrl+C stops at the next batch boundary.

visualize(e, outputs; <declared batch fields>) -> figure

Draw one sample. The signature follows metrics: the experiment and the outputs positionally, and the batch fields the method declares as keywords. It is called once per sample with the batch dimension dropped, so a rank-1 field yields its element: a Vector{String} of case ids arrives as a String. outputs is nothing in data mode (predictions = false), so one generic method can cover both modes, or two methods can split them:

visualize(::MyExp, ::Nothing; img, y) = data_panel(img, y)visualize(::MyExp, outputs; img, y)   = pred_panel(img, y, outputs)

There is no default method; a missing one is an error naming the experiment and the mode. Return anything save_figure can write.

with_progress_note(f, msg)

Run f with msg shown beside the current phase, and restore the enclosing note when it returns or throws. Nested blocks stack, and the innermost one is shown:

with_progress_note("loading index") do    for (i, f) in enumerate(files)        with_progress_note(() -> decode(f), "decoding $i/$(length(files))")    endend

A new phase or stretch clears every note, including those of blocks still open.

world_size(dist) -> Int

The number of processes. dist is always nothing today and world_size(nothing) == 1.

write_export(backend, model, ps, st, example_inputs; kwargs...) -> String

Write the artifact for backend and return the path written; the one method a backend implements. model is a callable model(inputs, ps, st) -> (outputs, st), the form a Lux-style tracer expects, that carries export_view(e), the routing, the preprocess and any echoes, and returns forward's full output. ps and st are host trees, st in test mode, and example_inputs is a tuple of host wire arrays at the first batch size. Keywords, all supplied by export_model: dir, name; input_names and output_names in program order; output_select, mapping model's output to the ordered tuple that ships; client_inputs, client_outputs, postprocess (each nothing or the hook's value); batch_sizes, each a separately compiled program; provenance, already merged.

@experiment struct MyExp ... end

Declare an experiment type. The macro generates the struct, a @kwdef-style keyword constructor, the device_fields and host_fields traits, the config_metadata table, a compile_view method, the Base.show methods, and a ?MyExp docstring.

@experiment struct MyExp    "Weight of the auxiliary heatmap loss relative to the primary term."    aux_weight::Device{Float32} = 0.25f0    "Number of decoder blocks. Structural: changes the compiled graph."    n_layers::GraphConst{Int} = 4    "Epochs to train for. Driver-only: never read inside a traced function."    max_epochs::Int = 40          # unmarked, so Host: the defaultend

A Device field is a traced input, a GraphConst field bakes into the graph and the cache key, and a Host field is invisible to the tracer and is the default. A macro is necessary because after setup a Device field holds a ConcretePJRTNumber and nothing at runtime says it was marked; every Device and Host field gets its own type parameter so what it holds can change, while GraphConst fields keep their declared type and the constructor converts to it. Docstrings are bare strings above each field, since comments never reach a macro.

The generated show prints values and never a device buffer, summarizing an array as its eltype and shape, because a Device{NamedTuple} of weights would otherwise print element by element. The macro is optional: a hand-written struct defines the three exported traits itself and may opt into the display with the Base.show one-liners the macro expands to (_show_experiment).

ReactantNitro.cache_stats() -> (; hits, misses, entries)

The compile cache counters. A test uses them to check that a fixed-LR loop trains a two-layer MLP for two epochs on CPU with no recompile after step 1.

ReactantNitro.cache_reset!() -> nothing

Empty the cache and its counters. Not part of the run path: it exists so a test can assert a miss, and so a session that has changed something the key cannot see (a const in the user's module, the cache's one documented hole) has an alternative to restarting the REPL.

ReactantNitro.check_compile_options(opts) -> NamedTuple

Validate the compile_options keyword of Nitro. It is normalized to the keywords every program's Reactant.Compiler.compile call receives: either a NamedTuple of Reactant's compile keywords (optimize, cudnn_hlo_optimize, transpose_propagate, xla_debug_options, ...) or a Reactant.CompileOptions, which becomes (; compile_options = opts). The keywords are exactly those Reactant.Compiler.compile accepts; the remaining CompileOptions fields (the pass switches) need a CompileOptions value.

Nitro(e; compile_options = (; xla_debug_options = (; xla_gpu_exhaustive_tiling_search = true)))Nitro(e; compile_options = (; cudnn_hlo_optimize = true, transpose_propagate = :down))Nitro(e; compile_options = Reactant.CompileOptions(; disable_slice_to_batch_passes = false))

An unknown keyword is refused here rather than minutes later at the first compile, as is an unknown field of xla_debug_options, xla_executable_build_options or xla_compile_options. So is donated_args other than :auto: the gradient accumulator and the parameters are written in place and must be donated, or every superseded buffer waits for the GC.

ReactantNitro.git_provenance(root) -> Dict{String,Any}

Repository state at root: git_commit, git_tree_sha1, git_branch, git_dirty, repo_remote, and on a dirty tree git_diff, the full git diff --binary HEAD, so the commit plus the patch reconstructs the exported code (untracked files are not captured). Also exported_at (UTC) and julia_version. A field git cannot answer is omitted; a root outside a work tree warns and omits every git field.

ReactantNitro.epoch_token(source)

An epoch counter that each begin_epoch! call increments by one, or nothing (the default) to skip the check. The fan-out asserts epoch_token(source) == before + 1 around its one begin_epoch! call, which catches a missing re-plan and a doubled one; a doubled one skips a generation of epoch-seeded augmentation. Define it as an Int field the source increments.

ReactantNitro.check_source_options(source, name::Symbol, cfg) -> nothing

Reject or warn about source options the resolved pipeline cannot honour. Called at setup for every split, :inline ones included, with the unwrapped source. The default does nothing; ReactantNitroMLUtilsExt implements it for MLUtils.DataLoader. cfg is the split's resolved prefetch_config, since an option implemented inside Base.iterate takes effect on the single-producer path and has no effect on the fan-out, which never calls iterate.

ReactantNitro.release!(nitro) -> nothing

Release every split of the handle's data collection through the source trait's release! method, once per handle. train! does this itself on reaching Done or Failed; call it for a handle that only validated, predicted or exported, whose sources would otherwise stay open. A split whose release! throws is warned about, and the other splits are still released. The handle cannot iterate a released source again.

ReactantNitro.release!(source) -> nothing

Release what a source holds open, such as a data server, a client or a file. Optional; the default does nothing. The framework calls it once per handle on every split's source, unwrapped from its PrefetchIterator or NoPrefetch, at the first Terminal phase (Done or Failed) the handle publishes, after the phase monitors have run. For a handle that only evaluates or exports, and so never reaches a Terminal phase, call release!(nitro). An error a method raises is logged as a warning and does not stop the other splits.

A method must tolerate a second call, since a source may be released explicitly and then again at Terminal. A collection shared between two handles (Nitro(e2; data = nitro.data)) is released by whichever finishes first; give each handle its own collection.

The source's package defines this method with the rest of the trait. For a public loader the framework may ship the extension, as with ReactantNitroMLUtilsExt; otherwise the loader's package declares ReactantNitro as a weak dependency and ships the extension itself.

ReactantNitro.progress_reporter!(f) -> previous

Set the progress reporter and return the previous one. nothing disables progress output; the default is default_progress_reporter. f is called as f(verb::Symbol, label::String, total::Int, epoch::Int, max_epochs::Int), with verb one of:

• :begin, once per stretch of work, with its label ("train", or the split name), its total units (0 when unknown), and the epoch position.
• :step, once per completed unit. The other arguments are placeholders.
• :end, once when the stretch finishes, however it finishes.
• :done, once when the last stretch of an entry point is over, so a reporter reusing one terminal line can close it. label is the closing text, or "" for the default.
• :phase, when the current stretch starts or stops doing something that produces no units; label names it ("compiling gradient") or is "" when ordinary work resumes.
• :note, when the note user code set with progress_note! or with_progress_note changes; label is the note, "" when there is none. A new :phase or :begin clears it.

A reporter that throws is switched off with a warning, and the run continues.

ReactantNitro.default_progress_reporter(verb, label, total, epoch, max_epochs) -> nothing

The progress reporter installed at load. At each stretch's :begin it picks, in order: progress_bar_reporter when the session is interactive and stderr is a terminal; progress_log_reporter when the current logger accepts ProgressLogging's level, which is a notebook or a REPL with TerminalLoggers; nothing otherwise, which is a CI log or a captured transcript. The choice holds for the stretch. Install either reporter directly with progress_reporter! to force one.

ReactantNitro.progress_bar_reporter(verb, label, total, epoch, max_epochs) -> nothing

The built-in terminal reporter: one bar per entry point on one terminal line, filled by the run's fraction so the ETA is the run's, described as epoch 3/40: train [compiling gradient] and ending as done: 40/40 epochs. Draws only when the session is interactive and stderr is a terminal.

ReactantNitro.progress_log_reporter(verb, label, total, epoch, max_epochs) -> nothing

The reporter for an environment that renders log records: Pluto, VS Code, a REPL with TerminalLoggers. Emits ProgressLogging.Progress records under one id per entry point, opened at the first :begin and closed with done = true at :done, with the same name and fraction progress_bar_reporter draws. Emits regardless of whether anything is listening; the logger drops what it does not accept.

ReactantNitro.MetricHistory

The per-epoch metric series a handle keeps, as returned by history. Each trained epoch adds one row, (; epoch, step, train_loss, metrics...): step is the optimizer step the epoch ended on, train_loss the epoch's mean train loss (NaN if it ran no micro-batch), and the rest that epoch's finalized validation metrics as host values (none without a val split), under the names they were reported with, so the validation loss is loss. The train loss is the one column named with its split, since the table holds both. Integers and ranges select epochs, symbols select metrics, and a selection always keeps epoch and step:

h = history(nitro)h[7]                         # the row for epoch 7, a NamedTupleh[end]                       # the last epoch's rowh[10:20]                     # epochs 10 to 20, still a MetricHistoryh[:acc, :macro_recall]       # two metrics, every epochh[10:20, :acc]               # bothh[step = 5_000:20_000]       # the epochs whose closing step is in that rangeh.acc                        # the column as a Vector, for a plot or a thresholdh.epoch, h.step, h.train_loss  # the axes and the train loss

It is a Tables.jl table with column access, so DataFrame(h) and CSV.write(path, h) accept it; with a Makie backend loaded, plot(h) draws it (see history_series). The columns are epoch, step, train_loss, the checkpointer's selection metric if it has one, then the other scalar metrics in order of first appearance; an epoch lacking a metric has missing. A non-scalar metric, such as a confusion matrix, is not a table column, but propertynames(h) lists it and h.confusion returns it per epoch.

ReactantNitro.history_table(h, height, width) -> NamedTuple

Everything a renderer needs to draw h in a height x width terminal: title, labels, cells::Matrix{String}, alignment, best_row, gap_rows and note. Rows are thinned with thin_rows. Columns after epoch, step, train_loss and the selection metric are dropped from the right until the table fits width, and note says what was left out and how to select it.

ReactantNitro.history_series(h; x = :epoch, metrics = nothing) -> NamedTuple

Everything a plot of h draws, computed without a plotting package: title, xlabel, and panels, one per curve, each (; name, x, y, best, subtitle). x is :epoch or :step. metrics is nothing (the checkpointer's metric if the history has it, else every scalar metric but train_loss, else train_loss), :all (every scalar column including train_loss), one name, or a collection of names. best is the (x, y) point of the checkpointer's chosen epoch on that metric's panel, else nothing. A row lacking the metric is skipped; a non-finite value is kept, so the line breaks where the run did. The Makie extension turns this into plot(h).

ReactantNitro.thin_rows(n, best, budget) -> Vector{Int}

Which of n rows to show in budget lines: windows around the first, the best (or nothing) and the last row, grown one row at a time round-robin until the budget is spent, with 0 marking a gap. The best window grows in both directions, so it gets twice the context of the ends.

StrippedHost{name}

The value compile_view substitutes for a Host field, carrying the field name as a type parameter. Using one raises through _stripped_error, which names the field and the two fixes. Common operations (getproperty, convert, indexing, iteration, arithmetic, broadcast) have methods that raise it; other operations fall back to Base and fail with a less helpful error. A sentinel that is only stored, compared or passed along raises nothing.

ReactantNitro.publish_phase(nitro, phase; info...) -> nothing

Fire the handle's monitors for phase without recording it on the handle. Used for Repl and Starting, which describe the process rather than the run: recording Repl in nitro.phase would erase how the run ended. There is no repeat guard, since there is no stored phase to compare against; the depth counter limits this to one event per outermost call. A Terminal phase also releases the handle's data sources.

ReactantNitro.publish_phase(phase::Phase; info...) -> nothing

Fire the module-level monitors for phase while no handle exists yet, during Nitro(e), where build_data may start and compile a data server. step, epoch, info.nitro and info.logger are nothing. The handle later adopts the same functions, so one monitor sees the events in order.

ReactantNitro.work_in_flight() -> Bool

Whether any public entry point is executing in this process; false is the state Repl announces. An entry point that throws before it has a Nitro publishes no Repl and leaves the last published phase in place, so a monitor polling on an interval should check this as well.

ReactantNitro.set_phase!(nitro, phase; info...) -> nothing

Record phase on the handle and publish the transition; the only writer of nitro.phase. Setting the current phase again is a no-op. Extra keywords are merged into info, which is how the transition out of EvalStepping carries metrics. Entering a Terminal phase also releases the handle's data sources, after the monitors run.

ReactantNitro.graphconst_field_hash(ev) -> UInt

A hash of compile_view(e)'s GraphConst fields only. Both exclusions are needed: compile_view strips Host and leaves Device in place, so hashing the view would hash device scalars that change every step. Each field's instance is hashed, which is content-based only if its type has a hash method; a struct holding a Vector falls through to objectid, and assert_graphconst_hashable refuses that at setup.

ReactantNitro.assert_graphconst_hashable(e) -> nothing

Refuse at setup a GraphConst value that does not hash and compare by content, since the compile cache and the resume check both rest on hash and isequal meaning "the same configuration". The check: a value and its deepcopy must hash equal and be isequal. A GraphConst{Vector{Int}} passes, since arrays hash by content; a struct containing a mutable field fails, since Base.hash falls through to objectid for it. Such a value would cost a full recompile per handle and a spurious resume refusal with an identical-looking diff, and mutating it in place would reuse a program compiled for the old value. It is refused rather than given a structural hash, which would change every existing key and decide equality for a type whose author has not.

ReactantNitro.check_config_compatible(record, e, path; weights_only = false) -> nothing

Compare a record's GraphConst fields with e's. These are the fields the compile cache is keyed on, so for a resume a change means the run would train a different compiled program, and it refuses with a diff. A weights-only load rebuilds the program from e and reuses the parameters alone, so it warns with the same diff; the parameter tree is checked separately. A field the record predates is compatible while it holds its declared default, which by convention reproduces the behaviour from before the field existed. Device fields are recorded separately in devices, and Host fields may change on resume. Runs after derive has merged, so a derived field is compared under the kind it was declared as.

ReactantNitro.check_record_tree_compatible(record, layout::FlatLayout, path) -> nothing

Refuse a record whose parameter tree does not fit the model this experiment builds, naming the first mismatching leaves by keypath with both sizes. layout is the fresh model's, built before the record's parameters replace it. Keyed by keypath, so a changed param_group is not a mismatch.

ReactantNitro._stripped_error(::StrippedHost{name}, op)

The message when traced code reads an unmarked field. The framework cannot know which fix was meant, so it offers GraphConst (bake, recompile per value) and Device (traced input, no recompile), plus the commonest answer: read it host-side.

ReactantNitro.check_control_readback(v, what, name) -> Float64

The validated readback for scalars the framework branches on (the early-stopping metric and the checkpoint selection metric): returns Float64(v) and raises if v is not a finite real. A logged NaN costs a missing point on a chart; acting on one could stop a healthy run or keep a stalled one going.

ReactantNitro.with_io_retry(f; attempts::Int = 5, backoff::Real = 0.5) -> f()

Call f, retrying up to attempts times in total when it throws a transient I/O error (see is_transient_io). The wait before retry k is backoff * 2^(k - 1) seconds, and each retry logs a warning. Any other error, and the error from the last attempt, is rethrown. Checkpoint reads, writes and rotation deletes go through it, since a failed checkpoint write can lose hours of training; a write goes to a temporary path that is then renamed into place.

ReactantNitro.PrefetchStream

The multi-producer fan-out, held as one object so close_stream! can stop every task. Iterating it yields the same (host, device) pairs as the other paths.

coordinator  ->  Channel{Int}(workers)             the batch indices 1:n   N workers ->  Channel{Tuple{Int,Any}}(workers)  HOST batches, with the index that produced them      1 transfer task -> Channel{Tuple{Any,Any}}(device_batches)  (host, device) pairs         the training loop   credits  <-  Channel{Nothing}(host_batches)     one token per host batch, returned on emit

credits holds host_batches tokens; it is a Channel rather than a Semaphore so teardown can close it and wake a blocked coordinator. marks records delivery, one byte per batch index; a Vector{UInt8} rather than a BitVector because workers write it concurrently.

ReactantNitro.auto_prefetch(collection) -> collection

Wrap every split in a PrefetchIterator at the defaults, unless it is already a PrefetchIterator or a NoPrefetch. Eval splits are included, since a forward-only step leaves less device time to hide the data path behind. Called at setup after derive, the schema probe and the contract checks, so those see the collection as build_data returned it.

ReactantNitro.batch_stream(split, routing, mesh = nothing) -> stream

The training loop's batch source, yielding (host_batch, device_batch) pairs. The loop needs the host batch for the schema checks and a :host train_metrics, and the device batch for the gradient program. An exception in a producer is rethrown at the consumer, since the Channel is closed with it. The checks run on the consumer, against the host batch, so errors are raised on the training task.

Three paths: the fan-out (PrefetchStream) when workers > 1 and the source is fanout_capable; one producer feeding a Channel of device_batches otherwise; and an inline generator for a NoPrefetch split. All go through prepared_stream with a prepare function from host batch to device payload: a plain transfer here, pad then transfer in eval_stream.

ReactantNitro.eval_stream(split, xfer, batch_size, mesh) -> stream

batch_stream for an evaluation pass: a short final batch is padded to batch_size before the transfer, so on the prefetch paths the pad and the transfer overlap the previous batch's forward. xfer is the routing whose fields are transferred. The yielded host batch is unpadded, and the consumer recomputes n_real from it rather than carrying it through the stream.

ReactantNitro.check_batch_at(source; values = false) -> nothing

Assert that batch_at agrees with the source's own iterate, batch for batch. It is the only check that catches a batch_at using the wrong index units, and it needs no accelerator, so it belongs in a model package's test suite. values = false compares field names and sizes, which works for a stochastic loader; values = true also compares contents and needs a deterministic source, since the two passes may draw different plans.

ReactantNitro.close_stream!(stream) -> nothing

Stop a stream's producers and free the device batches still in its channel. Called from the consumer's finally, so an early exit (request_stop!, a non-finite loss, an error) does not leak a task and device_batches x batch of device memory. It closes each channel, which stops each stage, and does not wait on the tasks: a worker in the middle of a slow fetch could delay teardown for seconds, and only the transfer task holds device memory. A no-op on the inline path.

ReactantNitro.fanout_capable(source) -> Bool

Whether source defines both begin_epoch! and either form of batch_at. There is no generic three-argument forwarding method, since one would make hasmethod return true for every source.

ReactantNitro.prefetch_config(split) -> (; device_batches, host_batches, workers, ordered, path)

The prefetch settings a split resolves to, which may differ from the requested ones. path is:

• :fanout: workers producers over batch_at, delivered in the source's order.
• :fanout_unordered: the same with ordered = false, delivered as the workers finish.
• :single: one producer over the source's own iterate, because workers == 1 was asked for.
• :materialized: one producer over a Vector of already-built batches; nothing to spread.
• :single_no_trait: one producer because the source lacks the trait. The setup warning's case.
• :inline: no stream at all, a NoPrefetch split.

ReactantNitro.strip_rules(opt_state) -> states

Replace every Optimisers.Leaf in an optimizer state tree with its state, keeping the structure. This is the traced half of the manual-mode contract, rules in and states out; merge_rules reverses it.

ReactantNitro.merge_rules(combined, states) -> opt_state

Re-attach the rules of the opt_state passed to the closure to the states it returned, rebuilding each Leaf(input.rule, returned_state, input.frozen). A structure mismatch is a framework bug and fails here, outside the trace.

ReactantNitro.RunSummary

One run of a run directory, as runs reports it: id (the run id, or "legacy" for the entries of a manifest written before runs were recorded), checkpoints (how many the manifest keeps), epochs (theirs, ascending), metric and mode (the selection of the run's most recently written scored entry, nothing for a run with no score), best (the best score under that selection), metrics (every (metric, mode) pair the run's entries were scored under, with counts), written (the wall-clock time() of its latest write, 0.0 for legacy) and parent (the checkpoint a branched run was resumed from).

ReactantNitro.check_source(src, kw) -> nothing

The shape of a checkpoint source, checked before anything is built, for kw in :weights and :resume: nothing, a path, or run => checkpoint. The run is :latest, :all, a run id, or a Nitro; the checkpoint is :best, :latest, or :current (a Nitro's in-memory state). weights also takes a bare Nitro (the same as other => :current), and resume takes false and :auto. Refused with the form to use instead: bare :best and :latest (no run named), :all => :latest (that is :latest => :latest), :current without a Nitro, and for resume a bare Nitro and other => :current, a continuation from memory, which is not supported yet.

ReactantNitro.resolve_source(ckpt, run_dir, src, kw) -> path, `Nitro`, or nothing

The one resolver for a weights or resume source that check_source accepted: a path is itself, a Nitro and other => :current are the handle (the in-memory warm start), and every other run => checkpoint names a record, resolved to its file through the manifest by resolve_checkpoint. run is :latest, :all or a run id of run_dir read by this run's checkpointer, whose current metric decides :best; other::Nitro is that handle's own run, read in its own run directory through its own checkpointer, whatever else has written there since. Raises when a record cannot be found, rather than falling back to fresh weights. resume = :auto is resolved by resume_path, not here.

ReactantNitro.resume_path(resume, ckpt, run_dir) -> path or nothing

The checkpoint resume names, or nothing for a fresh start. false (and nothing) resume nothing; a path is itself; run => :latest | :best is resolved by resolve_source. :auto is exactly :latest => :latest, the newest record of the directory's most recent run, announced with an @info naming the run and epoch it continues, except that a directory with nothing to resume in it is a fresh run with a @warn rather than an error, so :auto can stay set in a script that both starts and restarts runs. resumed_run decides which run the resumed handle joins.

ReactantNitro.resolve_checkpoint(ckpt, run_dir, run, which; what) -> (; path, epoch, run)

Resolve a run => checkpoint source of weights or resume to a record, from the manifest of the checkpointer's directory (its dir, else run_dir):

• run is :latest (the directory's most recent run: the run of the entry written last, legacy entries counting as oldest), :all (every run, only with which = :best), or a run id ("legacy" for the unnamed run of a manifest written before runs were recorded).
• which is :latest (the run's newest retained record) or :best (the best of the run's entries by the checkpointer's CURRENT metric and mode; across every run for :all).

Raises instead of returning nothing, since a caller who asked for trained weights and got fresh ones would not notice. Every miss names the directory, the runs it holds and runs(run_dir). A :best over entries scored under another metric warns that they were not ranked (warn_mixed_metrics), and a run with no entry under the current metric raises, naming the ones it has. A checkpointer of the user's own knows nothing of runs and answers :latest => :best | :latest through its own selected_checkpoint and find_latest methods. what is how the messages name the request.

ReactantNitro.warn_mixed_metrics(entries, ckpt, what) -> nothing

Warn when some of entries were scored under a metric or mode other than the checkpointer's, listing each pair with its count and saying they are kept. Selection ranks only the current pair, so those checkpoints are invisible to :best and to retention. Called at setup of a training handle, over the entries its retention decides over, and by every :best lookup, over the entries it considered. The manifest records a metric's name only, so a metric whose definition changed under the same name is not detected.

ReactantNitro.find_latest(ckpt, run_dir; run = :recent) -> path or nothing

The path of one run's highest-epoch retained record in run_dir, read from the manifest: what :latest => :latest and resume = :auto resolve to. Under keep_latest = false that is the run's most recent improvement, not where training stopped. run is a run id, nothing for the unnamed run of a manifest written before runs were recorded, or :recent (the default), the run that wrote the directory's most recent entry. Returns nothing when there is no such record or its file is gone; resolve_checkpoint is the form that raises and says what the directory holds.

ReactantNitro.selected_checkpoint(ckpt, run_dir; run = :recent) -> NamedTuple or `nothing`

The best of one run's checkpoints by the checkpointer's current metric and mode, as (; path, epoch, metric, score, run), read from the manifest alone. run is as for find_latest, the most recent run by default; train! passes its own. Only entries scored under the current metric and mode are ranked; an entry with no recorded metric predates the field and is ranked. Returns nothing with no checkpointer, no directory, no scored entry, or a missing winning file. show and train! use this quiet form; :latest => :best goes through resolve_checkpoint, which raises and warns about mixed metrics.