API reference

ReactantNitro

A general Julia training framework on Reactant.jl + Lux.jl.

The central object is a Nitro: Nitro(e) runs the setup sequence and nothing else, so validate, evaluate, and predict work with no training anywhere in the process. train! is a separate verb. A run is what happens when you train! a Nitro, which is why the run_* accessors keep that name while the object does not.

CheckpointRecord

The single authority on what a checkpoint holds. The record on disk is the snapshot handed to save_checkpoint! plus the two framework-stamped fields at the top.

opt_state is stored as host values and re-normalized to device residency on restore. Serializing ConcretePJRTNumber leaves would tie a checkpoint to a device configuration and JLD2 has no reason to round-trip them faithfully, so the flow is uniform with exactly one normalization point per path: write host, read host, normalize on the way in. Skipping the restore-path normalization reacquires the frozen-step-counter bug in full, on the default path.

A metric-comparison gate can read its curves out of this record's metrics field, which is why such a gate needs no logger and the framework ships none.

Writing a checkpoint.

Compiling <: Phase

Supertype of the four compile phases, which are slow by design. Subscribe here rather than to a leaf so that adding a leaf stays non-breaking.

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

Weight decay and decay-toward-the-pretrained-weights as one rule, differing only in the anchor:

anchor === nothing is a type-level distinction, so the branch resolves at trace time: plain decay emits no subtraction ops and carries no parameter-sized zero buffer.

Decay is the only name exposed, and there is deliberately no WeightDecay alias, which would collide with Optimisers.WeightDecay on the same concept. Name collisions are usually tolerable because users qualify, but that argument holds for train! and loss, where the two meanings are unrelated; here the two names would mean the same concept with different implementations, which is the case where a collision genuinely misleads.

The field is no_decay_mask rather than mask because "mask" once meant two unrelated things: this per-parameter exclusion (norm affines, biases, 1-D parameters), which stays automatic, and a deleted per-group partition. Only this one survives.

Decoupled only: Decay goes after the base rule, matching AdamW, so it is decoupled from the adaptive per-parameter scaling but still scaled by the learning rate. lambda therefore arrives pre-multiplied by η_g(t), and an LR schedule modulates regularization strength. Decay before the base rule would feed the decay term through moment estimation, producing coupled L2, which the framework does not support. Decay anywhere but last in a chain is an error, checked at setup on the chain the factory returns.

The anchor is an explicit field, never captured at init. An init-captured anchor would grab whatever x is at optimizer setup, which on resume is the restored weights.

Device{T}

Declaration-time marker for an @experiment field, legal only in that position. It is never instantiated and never appears in the generated struct: 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 in-trace e.aux_weight is the device value and arithmetic works with no unwrapping. The one marker does three jobs:

1. Device placement. Scalars become ConcretePJRTNumber, arrays ConcretePJRTArray.
2. Cache-key exclusion. A traced input cannot affect the graph, so it is excluded from the compile cache key by construction.
3. Schedulability. Constant and scheduled are the same device slot; only the write cadence differs, and switching between them does not recompile.

Marking is opt-in. Forgetting to mark costs a spurious recompile, which is visible and harmless; marking a structural field Device would silently lose constant folding or fail at trace time. The cost of over-marking is the lost constant folding: a Device threshold of zero emits the ops a plain host 0f0 folds away. Mark what you intend to revise or schedule.

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

See also Host, which is the default, and GraphConst, which bakes, and @experiment.

The run finished, whether by reaching max_epochs or by stopping early.

EarlyStopping(; metric = :val_loss, mode = :min, patience = 5, min_delta = 1f-4)

A small dedicated component, independent of the checkpointer even though both track improvement. The metrics can legitimately differ (checkpoint on mae, stop on val_loss), and the jobs differ (retention versus control flow); they share only the convention of a metric name and a :min/:max mode. Duplicating three lines of comparison beats the abstraction that would avoid it.

train!(e; early_stop = EarlyStopping(; metric = :val_loss, mode = :min,                                       patience = 5, min_delta = 1f-4))

:val_loss is not a magic name: it is what the default metrics emits. An experiment defining its own metrics names one of its own keys instead, and a metric no metrics emits is a setup error naming the available ones, checked before training rather than at the first epoch.

patience counts epochs without improvement; min_delta is absolute. Both match Keras and Lightning.

early_stop defaults to nothing, so a bare experiment does not early-stop. A framework that truncates your run by default is surprising, any patience value would be a guess, and the licence this framework takes to be opinionated is scoped to where the evidence is one-sided, which here it is not. The test suite is worded to match and asserts a bare experiment does not early-stop.

Stopping is graceful: finish the epoch, validate, checkpoint, finalize the logger, exit through the normal Done path. Aborting mid-epoch skips exactly the steps that make the run useful.

Stopping is recorded as stop_reason in the checkpoint record, since "completed 40/40" and "stopped at 37 on patience" are different outcomes. That also makes resume comprehensible: a run that early-stopped and is resumed with :auto would immediately re-satisfy the condition, and with the reason stored it says so instead of exiting silently.

request_stop! sets the same flag imperatively, from a REPL or a phase monitor. Both are checked once per epoch after validation.

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

Running the validation, testing, or inference loop.

ReactantNitro.ExportBackend

The supertype of an export target. One verb, write_export, dispatches on the backend value, exactly as the logging contract dispatches on the logger object.

The framework ships no method for any backend. ReactantServerBundle is a package extension on a weak dependency, which is the same care the logging contract takes with backends and the optimizer takes with schedules: this package depends on no logger, no schedule library, no batching library, no plotting package, and export is not the exception that breaks the pattern.

ExportCompiling <: Compiling

Tracing and compiling the export program, published once per export_model call rather than per compile-cache miss. Export retraces by design and never touches the compile cache, so there is no cache miss to key a transition on: the phase wraps the backend's write_export, which is where the one CPU compile per batch size happens, and the previous phase is restored when the bundle is written (or the call fails, since export publishes no Failed of its own).

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

One tensor's declaration, in the framework's own vocabulary rather than any backend's. A backend translates it; nothing outside a backend extension should know what it translates to.

shape is the Julia shape and batch_axis is a 1-based Julia axis, both matching every other shape statement in this framework. A backend that wants row-major network axes converts, which is the sort of thing having a framework type at all is for.

Which fields are required depends on which hook returns it, and the split is not arbitrary: the framework derives everything it can and asks only for what it cannot.

The two rejections are not tidiness. The executable side of a bundle is traced, so its shapes are facts rather than declarations, and the batch-last rule is one the framework already enforces on every array leaf of a batch and of forward's return. The client side is the output of a model.jl the framework never runs, so there it declares nothing and asks for everything.

axis_letters gives the manifest meaningful axis names, one Char per non-batch axis in shape order, e.g. ['w', 'h'] for an image. A -1 in shape marks a variable non-batch axis, which is how a postprocess with a data-dependent output width (a detector's detection count) is declared.

from names the SOURCE of an output, separately from what the bundle calls it. It is meaningful only on export_outputs and is refused elsewhere. See that hook for what it is for.

The run raised.

Tracing and compiling the gradient program.

GraphConst{T}

Declaration-time marker for an @experiment field, legal only in that position. It is never instantiated and never appears in the generated struct: the macro consumes it and records the field, which the framework reaches as setdiff(fieldnames, device_fields, host_fields).

A GraphConst field bakes as a trace-time constant and enters the compile cache key. The tracer sees the real value, so e.n_layers is a Julia Int with ordinary control-flow semantics; changing it is a different program, which is exactly why it is hashed and why a changed value recompiles rather than silently running the old graph. This is the field category for structure: anything whose value changes the shape or meaning of the emitted graph, such as layer counts, widths, or a seed that is deliberately meant to be part of the program.

Marking is opt-in and the opposite of the default. An unmarked field is Host, so a field that would previously have been left plain must now be written GraphConst to keep baking. The cost of over-marking is a recompile per distinct value; the cost of under-marking is a loud StrippedHost error when traced code reads the field, never a silent wrong program.

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

See also Device and Host and @experiment.

Host{T}

Declaration-time marker for an @experiment field, legal only in that position. It is never instantiated and never appears in the generated struct: the macro consumes it and records the field in host_fields.

A Host field is never converted, never hashed into the compile cache key, and never visible to the tracer. It does two jobs:

1. Keeps driver knobs out of the cache key, so changing max_epochs from 40 to 41 does not force a full recompile.
2. Keeps dataset-sized state away from the tracer. Reactant and Enzyme traverse the whole Const(e) argument while tracing, so anything dataset-sized reachable from the experiment (a sampler, an in-memory table, a materialized index) is walked element-wise, on one thread, every time a program is compiled. The cost is O(n_train) and it does not change the emitted graph, so nothing about the trained model looks wrong; it just gets slower the more data you have.

Host is the default: an unmarked field is Host. In practice the vast majority of experiment fields are driver knobs or dataset-sized state, so the marker is optional and max_epochs::Int = 40 means the same as max_epochs::Host{Int} = 40. Write the marker explicitly only where the field's host-ness would otherwise be surprising.

The guard is compile_view, which replaces every Host field with a StrippedHost sentinel. The framework passes that view to every trace site and uses the real e everywhere outside the trace.

Typical Host fields: max_epochs, output paths, log cadence, checkpoint retention, early-stopping patience, and any materialized dataset an experiment chooses to carry.

See also Device and GraphConst and @experiment.

JSONLogger(path = nothing) -> JSONLogger

The framework's default logger: writes JSON Lines (one JSON object per line) to path, defaulting to the run's metrics.jsonl. It is the value logger(e) returns when the experiment declares no logger, and setup pins a pathless one to the run's resolved run_dir, so a bare Nitro(e) writes runs/<Exp>/metrics.jsonl beside its checkpoints.

Every line carries a "type" discriminator:

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

Append mode, always, so a resumed run extends the same file rather than truncating it; that is also why logger_state returns nothing and no reattach! is defined: a ten-line file logger defines neither and keeps working.

logger = nothing is the documented opt-out, explicit per run or as an experiment field, and stays the public "no logging" value with its ::Nothing no-op methods.

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

The run's materialized state, and the public constructor for it: Nitro(e) executes the setup sequence and nothing else, so validate, evaluate, and predict all work with no training anywhere in the process. train!(e) is sugar for train!(Nitro(e)).

Every keyword belongs here and train!(nitro) has none. This constructor is the single authority for the keyword defaults; train!(e; kwargs...) forwards everything to it.

Three constructions cover the cases with no training in them:

nitro = Nitro(e)                                  # fresh weights from build_modelnitro = Nitro(e; checkpoint = "runs/x/latest")    # trained weights, no training this processnitro = Nitro(e; data = (; test = loader))        # supply data directly, skip build_data

Opaque, with accessors. Users never construct or mutate one field by field; read it through experiment, parameters, states, run_dir, current_step, current_epoch, phase, and binding_report, and write to it through request_stop!.

On the name. The object is a Nitro; a run is what happens when you train! one. That is why run_dir, run_id, run_url, run_ref, and the "run phases" all keep run: every one of them names the event. Trainer would be the conventional choice and is the wrong one, since this object is constructed for evaluation and serving with no training anywhere in the process.

What it deliberately does not hold is the Enzyme shadow dps, which is allocated inside the gradient program on every invocation and never crosses a boundary. A dps field here is the natural way to write that bug.

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

The recorded preset form: build the experiment from from_preset(E, preset) and record which named recipe this run claimed, so the name reaches the checkpoint record and log_params!.

n = Nitro(MyExp, :baseline; run_dir = "runs/baseline")

Keywords are split by which set they belong to, in one rule: a Nitro keyword goes to Nitro; anything else must be a field of E and goes to the recipe. Anything in neither is an error naming both valid sets.

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

A name that is both keeps going to Nitro, which preserves the collision resolution exactly: max_epochs, seed, run_dir, accum, n_devs, and gradient_clip_norm are struct fields, legal preset keys, AND run keywords, and here they are unambiguously the run keyword because this is a Nitro call. The preset acts at construction and the run keyword acts at run level, which is the layering. Nothing that resolved to the run level before this split moves.

The keyword set is derived from this constructor's own declaration, so it cannot go stale as keywords are added.

Overriding a field used to require going through from_preset and naming the preset a second time:

n = Nitro(from_preset(MyExp, :baseline; aug_rotate_deg = 9.0); preset = :baseline)

Prefer the form at the top. That one still works and is what this is shorthand for, but it names the preset twice and nothing checks the two agree, so passing preset = :variant there records a name the experiment was never built from, silently, and the recorded name is the whole point of recording one. Reported from the first real use, where every override the model needed was a field rather than a keyword, which made the two-name form the common path instead of the escape hatch.

Recording a name for a modified experiment is still accepted, deliberately: a preset's values are struct fields by the time anything sees them, so per-value provenance would be a claim the type system cannot back. What the split removes is having to say the name twice to get it.

ReactantNitro.NoPrefetch(source)

The only way to decline prefetch, and it is deliberately a marker rather than a number.

The framework wraps the train split in a PrefetchIterator at its own defaults, so a split that must run its host data path inline on the training task says so with this. It is a visible, greppable declaration that a reviewer will question; a numeric prefetch_device_batches = 0 field on an experiment reads as ordinary tuning, which is precisely how one model ran its entire data path inline for weeks with nothing contradicting it.

iterate is a passthrough, exactly as PrefetchIterator's is, so this is a valid data source anywhere one is expected.

Tracing and compiling the optimizer program.

Phase

Root of the run-phase tree. The framework publishes phase transitions through register_phase_monitor! so an external heartbeat, watchdog, progress display, or dashboard can be written without the framework shipping one. The registry is named for what gets built on it: a monitor is anything that watches a run and never steers it.

This is framework-level because phases differ in duration by four orders of magnitude: a compile takes hundreds of seconds while a training step takes milliseconds, so a watchdog with a fixed timeout kills runs mid-compile. The Compiling supertype is the general part of that signal; expected durations are site policy and no timeout table ships here.

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

Prefetch: workers producer tasks building host batches, one transfer task performing the H2D copy, and a Channel of device_batches (host, device) pairs feeding the training loop. This is what keeps the device fed across variable host latency.

The two buffer knobs are one per side of the transfer, and each names what it costs. device_batches is how many batches sit ON THE DEVICE staged ahead, and host_batches is how many may exist ON THE HOST at once: built but not yet transferred, counting the workers' hands, the hand-off channel, and the reorder buffer together. Neither is per-worker.

The framework wraps the train split with this automatically (see auto_prefetch), at these defaults, so a user normally never writes it. It stays public for the case where the defaults are wrong, and NoPrefetch is how a split declines entirely.

• workers is the knob that matters. Staging is lookahead, not concurrency: one producer at 1.5 s/batch cannot feed a 0.4 s consumer at any device_batches, because the buffer simply stays empty. A worker count is what turns host + device per batch into max(host / workers, device), and it requires batch_at and begin_epoch! on the source. Without them this falls back to one producer and setup says so.
• Resident device memory is (device_batches + 2) x batch: device_batches in the channel, one in the transfer task's hand, one in the consumer's. That + 2 is why the default is 1.
• Resident host memory is about host_batches x batch, and the default of 2 x workers is what keeps every worker able to build one batch while another waits to be transferred.
• The two bound different failures. device_batches back-pressures a SLOW CONSUMER: when the device channel fills, the transfer task blocks, stops draining the hand-off, and the workers block behind it. host_batches back-pressures a SLOW BATCH, which the device side cannot see at all: with ordered delivery a straggler leaves the device starving while finished batches pile up behind it, so the ceiling has to sit upstream at the coordinator.
• The transfer is on ONE task, not on the workers. Two reasons, and the first is fatal to the alternative: bounding device memory would need a semaphore released when the consumer is done with a batch, and there is no such moment (XLA execution is asynchronous and the executable holds its inputs, which is why free_batch! refuses to free a consumed batch). Second, it keeps every PJRT buffer creation on one task, exactly as the single-producer path did.
• Buffer lifetime is a known sharp edge: freeing must be explicit rather than left to the GC. The eval loop applies the same principle to validation outputs.
• Failure and cleanup. A worker that throws must surface at the consumer rather than hanging it, and early exit must stop every task and free what the stream holds. This path only runs when something has already gone wrong, so it is tested by deliberately throwing mid-epoch.
• ordered = true is the default, and it is what makes fanning out free. The workers finish out of order, so the transfer stage holds finished batches in a reorder buffer and emits them in the source's own order. A fixed seed therefore reproduces a run bitwise whatever workers is, and adopting batch_at on a source changes its throughput and nothing else.
• ordered = false trades that for throughput, which is the prior framework's behaviour and this one's before the reorder buffer existed. Emission then never waits on a straggler, at the cost that reordering repartitions the epoch into different accumulation groups, so a fixed seed no longer reproduces bitwise. It stays statistically equivalent to a different shuffle: every sample is seen exactly once, and check_train_divisibility's invariant (a group never spans an epoch boundary) is preserved either way.
• host_batches is enforced upstream, at the coordinator. A reorder buffer that simply drained the workers would grow to a whole epoch behind one slow batch, so the coordinator takes a credit before handing out a job and the transfer stage returns one after emitting a batch. host_batches may not be smaller than workers, and the constructor says so: the batch ordered delivery is waiting for must be inside the window, or a straggler deadlocks instead of stalling.

Batching, shuffling, and splitting are not shipped: MLUtils.jl covers those.

It is a DECLARATION, and the framework realizes it

The user constructs this in build_data, which setup runs early. The H2D transfer needs the resolved batch routing and the device placement, neither of which exists yet, so this wrapper cannot carry a producer: it carries the source and the settings, and the training loop builds the pipeline once the transfer it is supposed to overlap is defined.

The consequence for the user is nil, and the consequence for a reader of this file is that Base.iterate here is a passthrough. Iterating a PrefetchIterator by hand yields the underlying source's batches, in order, with no task and no channel, so it is a valid data source anywhere one is expected and cannot leak a producer when something other than the training loop iterates it. That is deliberate: an iterate that spawned a task would have nowhere to put the teardown, which is the exact leak this file warns about. It is also what keeps render and predict(nitro, loader) unaffected by the auto-wrap.

Prefetch applies to every split, and auto_prefetch wraps them all. The eval splits go through eval_stream, which pads a short final batch inside the producer so that the pad and the transfer both overlap the previous batch's forward.

ReactantNitro.ReactantServerBundle()

The StableHLO bundle target consumed by ReactantServer. The type is declared here so it can be named and exported; its write_export method lives in the ReactantServerExport extension, so using ReactantServerExport is what makes an export actually happen.

Repl <: Phase

The caller has control and no framework work is in flight: a REPL prompt between calls, or the moment after any public entry point returns. It is the phase a long-lived process spends most of its wall clock in, and the only one the framework publishes on the way out of its own code rather than on the way in.

It exists because a monitor cannot otherwise tell "finished, waiting for me" from "still working". Without it the last thing a monitor sees after a run is Terminal, which is indistinguishable from a process wedged during teardown, so a watchdog either kills a healthy idle session or waits forever on a dead one. What "too long to sit idle" means is site policy, exactly as for every other phase, so no budget ships here.

It is PUBLISHED but never RECORDED, the one place the framework separates those. Repl is a property of the PROCESS rather than of the run, so phase keeps answering how the run ended: phase(nitro) isa Done after a successful train! and isa Failed after one that raised. Writing Repl into that field would erase the outcome in order to record something that was never about the run, and a freshly constructed handle reports Starting rather than Repl for the same reason. publish_phase is the verb, and is worth reading for where the split is drawn.

Published only by the OUTERMOST entry point. train! calls validate once per epoch and render calls predict per batch; if an inner return published Repl, every epoch boundary would announce an idle session in the middle of a run, and a monitor that widens its patience while idle would stop enforcing the per-step budget for the rest of the run. A process-level depth counter is what prevents that; work_in_flight is the same counter, readable.

Setup is running: the setup sequence, before the first compile.

Stepping <: Phase

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

Terminal <: Phase

Supertype of the two end states. p isa Terminal is the one query a monitor needs in order to close itself.

TopKCheckpointer(; k = 3, metric = :val_loss, mode = :min, dir = nothing, name = nothing)

The default checkpointer, writing into run_dir.

It is constructible for an arbitrary experiment only because the default metrics guarantees :val_loss exists; an earlier draft left metric with no default and therefore had a batteries-included default that could not be built. If a user's own metrics does not emit the configured metric, that is a setup error naming the metric and the ones actually emitted, checked before training rather than at the first checkpoint.

Retain a latest alongside top-K, as a retention rule rather than a second file: never rotate out the newest, whatever it scored. On disk this is the union, K+1 files when the newest is not among the best and exactly K when it is. Resuming from the best checkpoint is not resuming from where you were.

Do not symlink latest into the top-K set. Rotation deletes the target when a better checkpoint arrives, so the link dangles through entirely normal operation with no user action, and checkpoint directories get copied between paths where tar/rsync/scp handle symlinks inconsistently.

Write a small manifest (file, epoch, metric, stop_reason) alongside; top-K bookkeeping needs it anyway and it lets resume find the newest without reading every file.

Passing checkpointer = nothing disables checkpointing and is the documented opt-out, with its ::Nothing no-op methods. TopK is the default rather than nothing because resume = :auto only works if something wrote a latest.

"Every N epochs", "best only", and "write to object storage" are ten-line implementations. The serialization format stays concrete (JLD2); swapping it means replacing save and load together, which is the only safe granularity.

name is the checkpoint filename, and nothing means adopt the experiment's checkpoint_filename at setup, exactly as dir adopts run_dir. Passing one pins it and setup never overwrites it. It is called with keywords, name(; epoch, step, metric, score).

Running training steps.

accum(e) -> Int

Micro-batches per optimizer step. Default _field(e, :accum, 1); also a Nitro keyword.

Unlike the other run accessors this one bakes: it reaches the gradient program as inv_n = 1/accum, a trace-time host constant, so it is in the cache key through this accessor's primary_world exactly as gradient_clip_norm is. A GraphConst field of this name is therefore correct rather than a trap, because it genuinely does change the compiled program.

backend(lgr) -> Any

Reach the native backend object. Default is the identity, because the default is that no unwrapping is needed: the logger a user passes IS their backend object. Pass an experiment tracker's own run handle, the extension supplies log_metrics! for it, and every other client function in that package stays callable on the same handle.

backend(w::MyWrapper) = w.exp

This escape hatch is what makes a small interface defensible rather than limiting, so the two decisions stand or fall together: log_image!, log_artifact!, log_curve!, and log_text! are dropped precisely because anyone needing them reaches the backend directly. Monitors reach the logger through info.logger, so logging something custom at validation time needs no dedicated hook.

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

One reverse-mode pass over an objective, w.r.t. one parameter subtree, for use inside a train_step closure. f(ps_sub, consts...) returns the scalar loss; the returned grads is a tree matching ps_sub, and everything in consts... is Const, so no gradient ever flows into it.

Every traced value the objective reads must be an argument, never a closure capture. Measured: a captured parameter subtree silently zeroed the gradient, with the primal loss correct, which is the worst failure mode, no error at any point. A plain struct like the model, which holds no arrays, is safe to capture; parameters, batches, and state go in consts....

The non-saturating GAN formulation falls out of the activities. In the generator's backward, the discriminator's parameters are a Const argument: the loss's adjoint flows through the discriminator's forward as a function of the fake data and never into the discriminator's parameters. No explicit stop-gradient op is needed.

The mechanism is the framework's own objective-wrapper pattern, verified by the same measurement as objective_wrapper: the objective is wrapped in a named function returning a one-tuple, the return activity is Enzyme.Duplicated rather than Active, and the wrapper is passed Const with the objective as a Const argument. dps = Enzyme.make_zero(ps_sub) is allocated inside this call on every invocation, never hoisted, so gradients never accumulate across calls, which is the same discipline the gradient program follows.

The whole train_step body is one traced program, so backward is compiled as part of it; the host-side analog is exactly what the framework's own tests use as their reference.

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

Optional, and half of the index-addressable opt-in. Produce batch i of the current epoch, independently of every other i, from a task that may be one of many running concurrently.

Two shapes, and which one you write depends on where the epoch's plan lives. A source that stores its own plan writes the two-argument form; a source that is an immutable declaration of a dataset, with nowhere to store one, has begin_epoch! return the plan and writes the three-argument form. MLUtils.DataLoader is the second kind, and the extension that supports it is two methods and no state.

Leave the third argument untyped. The capability check asks whether a three-argument method accepts any plan at all, so annotating it with the plan's concrete type hides it and the source falls back to one producer.

i is a BATCH index in 1:length(source), not a sample offset. That distinction is the single easiest way to corrupt a run with this trait: a loader whose own producer takes a sample offset needs the multiplication, and

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 compiles, runs, and trains an epoch on length(dl) heavily overlapping windows of the first few hundred samples, with a loss curve that still falls. No runtime assertion can see it, because a permuted or overlapping index set produces perfectly well-shaped batches; check_batch_at is what catches it, and it needs no accelerator.

Requirements on the implementation:

• Thread-safe with respect to the source's shared state. Per-task scratch (a shared-memory segment, an RNG, a mutable buffer) belongs in a TaskLocalValue or is allocated per call.
• A pure function of i and the epoch's plan. Anything that re-plans the epoch belongs in begin_epoch!, which the framework calls exactly once before any batch_at. The three-argument form makes that literal: the plan arrives as an argument, so N workers read one immutable object rather than racing on the source's fields.
• Never nothing for i in 1:length(source). A loader whose sequential producer returns nothing at exhaustion must not forward that; the framework raises on it.

Define this and begin_epoch! to opt in. Defining only one is a deliberate dead end (see begin_epoch! for why).

ReactantNitro.begin_epoch!(source) -> plan

Optional, and the other half of the index-addressable opt-in. Re-plan the epoch. Called exactly once per epoch, on the training task, before any job is dispatched to any worker.

Return nothing if the source stores its own plan, which is what every stateful loader does and what this returned before there was anything to return; the framework then calls the two-argument batch_at. Return the plan itself if the source cannot store one, and the framework hands it back to every three-argument batch_at call for that epoch. The returned value is opaque: the framework holds it, passes it along, and drops it when the epoch ends.

This is what Base.iterate's initialization used to do, and the reason it cannot stay there: the fan-out never calls iterate on the source at all, so a source that re-plans in its iteration init and exposes batch_at would train every epoch after the first on epoch 1's plan, silently and with a plausible loss curve.

That is why both methods are required to opt in, rather than batch_at alone with a no-op default here. A no-op default would make the stale-plan bug the default outcome for exactly the loaders that most need the fan-out. With both required, a source that supplies only batch_at gets the single-producer path and a warning naming this function, which is a loud opt-out instead of a silent corruption.

A source that genuinely needs no re-plan opts in with a one-liner:

ReactantNitro.begin_epoch!(::MySource) = nothing

A source that is a declaration rather than a cursor returns its plan instead, and needs no mutable field at all:

ReactantNitro.begin_epoch!(d::MyLoader) = shuffled_index(d.rng, d.n)ReactantNitro.batch_at(d::MyLoader, i::Integer, plan) = _build_batch(d, plan, i)

Implement it once and call it from Base.iterate too, so the two entry points cannot drift:

Base.iterate(s::MySource) = (begin_epoch!(s); _first_batch(s))

length(source) is read after this returns, so a source whose batch count changes with the plan is handled correctly.

Setup calls it once, before any epoch. Setup draws one batch through first(source) to learn the batch schema, and that goes through Base.iterate, so a source counting its own epochs sees N + 1 calls across N training epochs. That is pre-existing behavior rather than something this trait introduced, since the probe always called iterate, and it is why a loader that plans its epoch eagerly in build_data needs a flag saying the first plan is already drawn: the probe is what consumes it.

binding_report(nitro) -> String

The binding report as plain text: where every configured value actually bound. It is a diagnostic, not a check; it computes nothing the run does not already compute and it never fails. Setup hands this same text to log_other!(lgr, "binding_report", str), so the run's record carries it whether or not anything displayed it.

What you read is normally show(nitro), which appends the report's sections to the handle's own, so one table answers both what the run holds and where each value came from. This accessor is for the text: a log line, a file, a diff between two runs.

It exists because the rules that resolve a learning rate, a schedule key, and a per-group accessor are individually simple and jointly hard to hold in your head, and because a wrong binding is otherwise invisible until a curve looks strange three hours in.

build_data(e, dist) -> NamedTuple

Required. Return the run's named data collection, (; train, val) or (; train, val, test). Named and extensible, so calling evaluate with no test supplied is a clear error rather than a positional mistake.

A data source is anything iterable that yields concrete NamedTuples of host arrays, and supports length. A Vector of batches, an MLUtils.DataLoader, or a bespoke sampler all qualify. Loaders yield host arrays; the framework transfers each batch, applies batch-dimension sharding, and owns prefetching.

Use MLUtils.DataLoader unless you have a reason not to. Any source works, and the framework plays no favourites at runtime, but a DataLoader is the one shape whose mistakes are caught before they cost anything. Three things follow from it that nothing else gets for free:

• Its settings are checked at setup, from its own fields. A train split that keeps its partial final batch is refused before the first compile, rather than on the last batch of epoch one after two compiles and a full epoch of stepping. An eval split that drops one is warned about, because it silently shrinks the set every metric is computed over.
• It fans out with no effort. The MLUtils extension implements the index-addressable trait for it, so its host data path uses every thread, in the source's order, while a source that implements neither half runs one producer and says so.
• Batching, shuffling and collation are its job, and the framework deliberately ships none of them. shuffle = true reshuffles every epoch, which a Vector of batches built once in this function cannot do at all.

A bespoke sampler is entirely supported and sometimes necessary; it just puts the three above back on you, and the first one is caught late rather than early.

Three requirements on the source, each with a reason:

• Restartable. Setup draws one batch to learn the schema and resolve routing, then discards it, and the loop iterates from scratch. A one-shot source such as a bare Channel would silently lose its first batch.
• length. Read once, at setup, to fix the schedule horizon. It counts batches.
• The train split drops its partial final batch (partial = false for MLUtils.DataLoader, drop_last = true elsewhere), and its batch count must divide by accum. Eval splits should NOT drop theirs: the framework pads and slices, and padding is safe there because testmode makes a standard model a per-sample function along the batch axis. A model that genuinely mixes across the batch in test mode is the one case where dropping on eval too is right, since it avoids padding rather than tolerating it.

There is no prepare_epoch! hook: for batch in loader calls iterate afresh each epoch, so a loader that reshuffles or regenerates does so in its own iteration initialization.

dist is nothing and nothing dispatches on it; it is the seam a distribution layer would use. Write it untyped.

build_data sees the pre-conversion experiment, where Device fields hold plain host values. It is skipped entirely when Nitro(e; data = ...) supplies the collection directly, which is what makes serving possible.

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

Required. Standard Lux.

Pretrained loading must happen inside build_model, because the framework captures w0 immediately after it returns; loading afterwards would anchor decay_anchor = :w0 to the random init, silently.

It sees the post-conversion experiment, so Device fields already hold device values.

For a model with variant branches, lift the selector to a type with dispatch_variant so the branches fold at trace time rather than leaving the return type a union of every variant:

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

Cost is one dynamic dispatch at a boundary crossed once per run. Each variant gets its own Julia specialization, so many variants multiply Julia compile time, which is a different budget from XLA compile time.

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

The checkpoint filename, and a user hook: define a method for your experiment type to name checkpoints differently.

epoch-0006-step-4500-val_loss=0.0416831.jld2epoch-0006-step-4500.jld2                          # a run with no `val` splitepoch-0043-step-32250-mae=-1.23457e-07.jld2

Reading a run directory should answer "which checkpoint is good" without opening JLD2. The score is in the manifest and in every record too, but both need deserializing to rank three files.

The epoch stays first and stays zero-padded, so lexicographic order is training order. The step is unpadded, because the epoch already supplies the ordering and a width would be a promise that breaks above its own digit count.

score === nothing omits the segment rather than filling it, which is what a run with no val split produces; the retention rule keeps such a checkpoint by being newest. No metric segment means no metric was computed, and a token such as no_score would reserve a name a real metric could take.

The value needs no sanitizing, and that is a consequence rather than a convention. The selection metric goes through check_control_readback before anything uses it, so score is either nothing or a FINITE Float64, and %g on a finite double is always within [0-9.eE+-]. A NaN score does not produce an odd filename; it stops the run. The metric NAME is sanitized, because it is a user-supplied Symbol and a filename is not.

A method must accept kwargs... unless it wants a later keyword addition to be a MethodError:

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

The framework calls this through TopKCheckpointer's name, which is bound to the experiment at setup, so save_checkpoint! keeps its four arguments and never sees e: the host-value discipline exists to keep device-resident leaves away from the function that serializes, and the experiment is the one object in a run that holds them by design. A consequence worth having: the binding dispatches at WRITE time, so a revised method takes effect on the next checkpoint of a live run, unlike a revised checkpointer accessor, which is called once at setup.

checkpoint_info(path) -> NamedTuple

What a checkpoint says about itself, with none of its weights. Which epoch, which step, which preset, which run, how it stopped, and the validation metrics that were current when it was written.

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

This is the ONLY thing you should open a checkpoint with when you are asking a question about it, and it exists because the alternative kept being reinvented, badly. A record is one JLD2 entry holding one object, four of whose fields are the whole parameter tree, so there is no way to read epoch without materializing all of it, and every session that tried built its own reader: JLD2.Group (does not exist), JLD2.names (not public, wrong method), keys on the record (no method), then getfield on a record read in a process without ReactantNitro loaded, where JLD2 hands back a ReconstructedMutable whose fields answer to getproperty and NOT to getfield. Then the one that actually costs something: println the record, and the parameter tree goes to stdout. CheckpointRecord's show now refuses that, and this function means nobody has to get close to it.

Read it in a process that has ReactantNitro loaded. That is not a nicety either: the workspace needs the type, or JLD2 reconstructs a look-alike and the field access rules change under you. This function is also where the two field-shape migrations live, so an older record answers the same questions as a new one.

For a whole run directory, do not call this in a loop. read_manifest answers "which checkpoints are here, at which epochs, with which scores" from the run's own small manifest and opens no record at all; go to a record only for the fields the manifest does not carry.

checkpointer(e)

The run's checkpointer. Default _field(e, :checkpointer, TopKCheckpointer()); also a Nitro keyword. nothing disables checkpointing and is the documented opt-out.

An experiment whose selection metric is one of its own metrics keys rather than :val_loss should say so here, so that every run of it selects the same way without the caller remembering.

compile_view(e) -> e_trace

The stripped view of an experiment that every trace site sees: each Host field is replaced with a StrippedHost sentinel and everything else is passed through unchanged. Since an unmarked field is Host, this strips every field that is neither Device nor GraphConst.

The framework calls this at every trace site and passes the result as Const. The real e is used everywhere outside the trace: build_data, derive, metrics finalization, checkpointing, and every driver decision.

@experiment generates a method per experiment type, so the reconstruction is type-stable and costs one struct copy. The generic fallback here covers a hand-written experiment that defines host_fields itself, and is an ordinary accessor, so an experiment whose layout the default does not suit overrides it directly.

Note it strips Host and leaves Device in place, which is why the compile cache key is computed over this view's GraphConst fields only rather than over the view.

Leaving Device in place is right for every trace the framework itself invokes and wrong for the one trace it hands away. See export_view for the frozen view export uses instead, and for why a Device field that reaches a serialized artifact makes it unservable.

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 one of :device, :host, :graphconst. type is the declared type, so a Device field reports the type inside the marker rather than the device type it holds after setup. doc is the field's docstring or nothing, and line is its declaration line, so a validation error can point at the source. default is ReactantNitro.NO_DEFAULT for a field declared without one.

This is the source of the logged hyperparameter table, and the input any tool that renders an experiment's configuration reads.

The fallback synthesizes a record for a hand-written experiment from fieldnames, fieldtype, device_fields, and host_fields, reporting doc = nothing, line = 0, and default = NO_DEFAULT throughout. It reports fieldtype, which for a hand-written Device field is the device type after setup rather than the declared one; a hand-written experiment that wants the declared types in its metadata defines this method itself.

current_epoch(nitro) -> Int

current_step(nitro) -> Int

The optimizer step, not the micro-batch. With accum = N the two differ by a factor of N, and confusing them shifts the whole learning-rate curve by N.

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

What this group's decay pulls toward. Default :zero, i.e. ordinary weight decay.

:w0 decays toward the parameters exactly as build_model returned them, before any training step or restore. The literature calls this L2-SP, for "L2 distance to Starting Point" (Li, Grandvalet, and Davoine, ICML 2018, arXiv:1802.01483); the name is recorded once so the technique stays findable, and decay_anchor = :w0 is used everywhere else because it says what it does.

An explicit array is the extension point: initializing from A while anchoring to B is decay_anchor returning B.

decay_anchor = :w0 on a randomly initialized group anchors it to a random point. That is why :w0 goes on the pretrained backbone and :zero on the fresh head, why :zero is the default, and why the framework cannot warn: it cannot detect which groups are pretrained.

Anchoring anywhere but :zero puts a per-group anchor_checksum in the checkpoint record and makes resume verify it and refuse on mismatch. The arrays themselves are not stored, which keeps a checkpoint from roughly doubling for a :w0-anchored backbone, and the accepted cost is that a run whose anchored parameters come from a nondeterministic source cannot be resumed at all.

ReactantNitro.default_no_decay(keypath, param) -> Bool

The framework's built-in decay exclusions: every 1-D parameter. In Lux that is biases and every normalization layer's affine scale and shift, which is why one predicate covers all three.

Structural rather than name-based, so it does not depend on a layer naming convention. Exported so that no_decay can be extended rather than replaced, which is the common case:

ReactantNitro.no_decay(::MyExp, ks, x) = default_no_decay(ks, x) || (:fc in ks)

derive(e, data) -> NamedTuple

Values that genuinely depend on the dataset. Default (;). The framework merges the result into the experiment:

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

It runs before device conversion, so it always returns host values, and it sees the pre-conversion experiment. A derived Device becomes a device value excluded from the cache key; a derived GraphConst field becomes a baked constant included in it.

Prefer constants for structural values; derive only numerics. Deriving a shape-determining value makes the compiled shape a function of the data, so a different split silently changes the program and the cache key. Use a constant with an explicit bounds check at data load instead.

Distinguish derived from merely computed-late. A value computable from config alone belongs in an ordinary computed default, not here.

Resume recomputes, it does not restore, so resuming against changed data picks up new values; they are recorded in the checkpoint record so a change is visible after the fact.

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. nitro.e holds device values after setup step 5, so reading a field directly hands back a ConcretePJRTNumber or ConcretePJRTArray; this transfers it back. The counterpart to set_device!, and the supported way for host-side code to read a value it is also writing.

Errors on a field that is not Device, because a Host or GraphConst field is already a host value and getfield is the right way to read it.

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 stopping rule. Default _field(e, :early_stop, nothing); also a Nitro keyword. nothing means no early stopping, which stays the default so that a bare experiment does not truncate its own run.

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

Run any named split from the data collection. Distinct from validate rather than a second name for it, and it errors on a split name build_data did not return, naming the available ones, which is the reason the data collection is a NamedTuple rather than a positional tuple.

Shares one compiled forward and one compiled metrics with validation and inference, so moving between them never recompiles.

Like every entry point, this runs its loop on a worker thread when one is available, and Ctrl+C stops the split at its next batch boundary and surfaces the interrupt.

experiment(nitro) -> e

The post-conversion experiment, which is what every hook except build_data and derive sees. Its Device fields hold device values.

export_client_inputs(e) -> Vector{ExportSpec}

Optional. The client-facing spec of the INPUTS, for the one thing that is not derivable about them.

Their dtypes and shapes are not that thing. The wire preprocess lives inside the traced graph, so the executable inputs already ARE the wire inputs, and a manifest that repeated them would be repeating itself; with no method here the server falls back to the executable specs, which is the same answer.

What is not derivable is axis_letters. The tracer derives the executable input specs itself and has no way to carry letters into them, so a model that wants "whcn" in its manifest rather than auto-allocated letters has to say so on the client side. That is the whole reason this hook exists, and it is worth having because an axis named w documents a wire contract in a way a does not.

Like export_client_outputs it requires export_postprocess, and for the same reason: client-facing specs are rejected by the server when no model.jl is present.

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) -> Vector{ExportSpec}

Optional. What a client receives after export_postprocess has run, when that differs from what the executable emits.

Only meaningful alongside a postprocess, and the framework refuses the combination that does not make sense: declaring client outputs with no model.jl writes a bundle that fails when the server loads it, so it fails here instead.

These specs are the one place ExportSpec's full expressiveness applies. The framework cannot derive them, because a postprocess is opaque Julia it never executes, so dtype and shape are required, batch_axis is unconstrained, and both axis_letters and -1 variable axes are available.

batch_axis being unconstrained here means unenforceable, not encouraged. Batch-last is the convention everywhere else and the verification simply cannot reach past the executable boundary. One model predates the rule and keeps its batch-middle client tensor; there should be no new ones. Since batch_axis defaults to last, any explicit value is either redundant or an exception, which makes both greppable and neither silent.

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 wire contract: what a client sends, in the order it sends it.

The order is load-bearing and is not a presentation choice. It fixes the positional order the program is traced with, the order of the tensor names in the bundle, and therefore the order a client must supply. Reordering this vector changes the wire contract of the next bundle, silently, so treat it the way you would treat a struct's field order in a serialized format.

Each spec needs a dtype and a shape, because the framework builds the example arrays it traces from them. The batch_axis must be the last axis: batch-last is the framework's rule everywhere, and export does not get an exemption from it. The size at the batch axis is a placeholder, overwritten by each entry 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, not a path, which is the single largest simplification this surface makes. The handle already carries the restored weights, the layer state, the experiment, the preset and the seed, so export never parses a checkpoint file and no model's export code contains the words "load" or "checkpoint" anywhere:

using ReactantServerExport                       # the extension that provides the backendnitro = Nitro(e; checkpoint = "runs/x/best.jld2", data = (;))export_model(nitro, ReactantServerBundle(); dir = "export_out", name = "my_model_v1")

data = (;) is not a workaround. Nitro(e) runs setup and nothing else, so an evaluation or serving construction never needs the training data, and export is the purest case of that: it reads weights and traces a graph.

batch_sizes is a list because each entry is a separately compiled program, traced and stored independently. It defaults to [1], and widening it costs compile time and artifact size in proportion.

Export is a CPU trace and asserts one device. A sharded program is not servable as a bundle, and the alternative to asserting it here is discovering it when something tries to load the result.

Provenance is assembled from four sources, and this is the precedence, lowest first, because somebody will need to know which one wins:

provenance_root is how repository state reaches a bundle. Without it the bundle carries no git commit, no tree hash and no working-tree patch, and it says so by omitting those keys rather than writing empty ones. That case is a legitimate choice and it is also the failure mode worth naming: an export with no root SUCCEEDS, every check passes, and the manifest looks complete while being unable to say which code produced the artifact. Pass the repository root and the backend collects the rest, patch included. export_provenance explains why the framework will not go looking for it by itself.

The model layer needs no argument at all, which is the point of it being a hook: export_provenance_extra is called with the experiment and with the checkpoint the handle actually restored from, so neither entry point can forget to pass it and neither can pass a path that disagrees with the loaded weights.

What happens, in order: check the hooks agree with each other and with the batch-last rule, build example wire arrays at the first batch size, run the program eagerly once to derive the output shapes and verify the batch-last rule the backend's own derivation depends on, resolve the provenance, and hand all of it to write_export. The backend call is published as the ExportCompiling phase, so a phase monitor sees the minutes-long trace rather than a silent stall; the previous phase is restored when the bundle is written, or when the call fails.

The trace is minutes-long, one CPU compile per batch size, so, like every entry point, this runs on a worker thread when one is available: the interactive thread's logger tasks and REPL keep running for the whole export. ^C aborts the wait; the export itself finishes in the background and completes the bundle, because a partial bundle is worse than a late one.

export_outputs(e) -> Vector{ExportSpec}

Required. Which leaves of forward's output tree ship, and what they are called.

forward returns what training needs, which is routinely more than a client wants: an ODE model's return carrying kinetic-energy rows is the case this hook exists for. Naming a subset is honest and needs no second program, and a separate export_forward would be a second thing to keep in agreement with the first.

Each name must be a key of the NamedTuple forward returns, or, when forward returns a bare array, this must be the single spec that names it. dtype and shape are derived from the trace; supplying them turns them into assertions, which is worth doing for an output whose shape you want a bundle to fail on rather than drift on.

export_outputs(e::CropClassifier) = [ExportSpec("region_logits")]

from, and why it exists

A bundle's tensor name and the leaf that produced it are two different things, and forcing them to be one made an export concern rewrite the trained program. Two cases need separating, both found by porting a real model rather than by design review:

• A rename. The wire contract calls a tensor state_out and forward calls that leaf g_pred. Without from the only way to ship it is for forward to return a second, aliased leaf, which recompiles the gradient program to serve a serving detail.
• An echo. A serve-time postprocess receives the program's OUTPUTS and never its inputs, so a postprocess that masks by valid or scales by scale_px needs those wire tensors back as outputs. Without from the only way was, again, to make forward return them.

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

An echo must say so. A spec whose name happens to match a wire input is refused unless it writes from explicitly, because turning a client's own tensor into a program output on the strength of a name collision would be a guess. The value echoed is the WIRE value, which is the tensor the client sent, not the preprocessed one.

export_postprocess(e) -> String

Optional. The source of the model.jl that ships in the bundle, as text.

The framework never runs it and has no opinion about its contents beyond writing it where the backend says it goes. It is the place for work that does not belong in a compiled graph: a softmax, a decode, an assembly of raw logits into whatever a client actually wants.

Returning a source string commits you to also implementing export_client_outputs or export_client_inputs when the postprocess changes what a client receives, and returning nothing (the default) commits you to implementing neither.

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

Optional, and TRACED. The seam between the wire and the batch, and the reason the exported graph is not the graph that trained.

It receives one positional argument per export_inputs spec, in that order, and returns the NamedTuple batch forward declares its keywords against. Whatever it does happens inside the compiled program, so a client sends bytes and the executable converts them.

With no method, the framework maps each wire input to a batch field of the same name, which is right whenever forward already declares the wire tensors directly.

It obeys every rule a traced hook obeys: no Host field reads, since the compile view strips them, no data-dependent control flow, nothing that is not a Reactant operation.

AND IT IS ALSO CALLED EAGERLY, ON HOST ARRAYS, WHICH IS THE PART THAT SURPRISES PEOPLE. Twice, in fact: once by export_model's verification probe and once by the backend's own shape discovery, both before anything is traced. So it has to work in both modes, and the single most common thing it will ever do does not:

Float32.(img)     # MethodError: no method matching Float32(::Reactant.TracedRNumber{UInt8})

An integer-to-float conversion has no traced broadcast method, so the wire conversion needs two methods, one per mode. This is not a nicety; every model doing the ordinary UInt8 to normalized Float32 conversion hits it on the first export.

_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}

What the framework knows about how these weights came to exist, and nothing else.

Provenance is the same question whatever the artifact is, which is why it is resolved here rather than in a backend: the preset name travelling with the bundle is what finally lets a served model answer "which recipe produced you", instead of that fact living in a launch script nobody kept.

It returns the flat config from config_params, the preset name recorded on the handle, this package's version, the seed, and the run directory.

It deliberately does not guess at repository state. A git commit, a tree hash and a working-tree patch are site policy, not framework knowledge, and a framework that shelled out to git would be asserting that the process's working directory is the model's repository. export_model takes a provenance dictionary that is merged on top, which is where that half belongs. A backend may stamp its own facts underneath, and its own version string is its own business.

It stamps checkpoint when the handle restored from one, naming the file the restore actually read. That is framework knowledge and not a guess: the path was handed to the constructor, or, under resume = :auto, resolved by the framework's own search, and Nitro.checkpoint_source retains whichever it was. A handle built from freshly initialized weights omits the key rather than carrying an empty one, so "no checkpoint" and "some checkpoint" are distinguishable in the manifest.

It stamps the TRAINING RUN's id and url as trained_run_id and trained_run_url, taken from the restored record and not from this handle's logger. That distinction is the point: a checkpoint = path construction gets a fresh logger, so the handle's own run_id names the process doing the exporting. A manifest carrying that would name an experiment holding an export trace and no training metrics. With the record's id in the manifest, everything else about the run, the training commit, the branch, the full logged hyperparameter set, is one link away rather than something a reader has to infer from a run directory's name.

One thing it still does not carry: the checkpoint's epoch and its metric. A weights-only restore deliberately zeroes the epoch counter rather than continuing it, so the handle's epoch is not the checkpoint's; the record's epoch, step and metrics are available at construction and simply are not retained. Stamping them is the same two lines as the run id above if it turns out to be wanted.

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

The MODEL's half of a bundle's provenance, merged on top of the framework's and the site's. The default is empty, so this is opt-in and an experiment that has nothing to add implements nothing.

export_provenance stamps what the framework knows and a backend stamps its own facts underneath. Neither can know what a consumer of this particular model needs in order to use it: which labeling variant produced the head, the class names in the labeler's order, the crop geometry a client has to reproduce, the contract strings that make a served tensor interpretable. None of that is inferable from the tensor shapes, and a class count alone routinely fails to identify a model, since two different labeling schemes can be the same width.

Why this is a hook rather than something the caller merges in. It was a hand-merged dictionary first, and that shape has one failure mode which is worse than an error: omit the argument and the export SUCCEEDS, the arity check passes, the manifest carries the framework's stamps, and the bundle is untraceable while looking complete. Documentation had to warn about remembering the argument, and more than one model independently wrote the same function under a name of its own, so nothing generic could call any of them. A hook the framework merges removes the argument from both entry points at once, so the tool path and the hand-written path get the same bundle.

checkpoint is the file the handle's weights came from, or nothing for freshly initialized weights, and export_model passes Nitro.checkpoint_source rather than asking the caller to name it again. The nothing case is a discriminator and not merely an absence: a model whose provenance asserts anything about what it was trained against should report those fields as unverifiable rather than as verified when there is no checkpoint behind them.

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

Keys collide with the framework's and the site's at the model's own risk: export_model states the precedence and this half wins over both. The explicit provenance argument still wins over this, so a human can always override a hook.

export_view(e) -> e_export

The view of an experiment a frozen trace sees: compile_view's stripping of every Host field, and then every device-resident value that survives it read back to the host.

This is deliberately not the view a training step sees. A Device field is a traced INPUT, which is exactly right while the framework owns both ends of the call: it supplies the value at every invocation and changing the value costs no recompile. An exported bundle has no such owner. It is executable(inputs..., weights...) and nothing more, so a value that is neither a declared input nor a serialized weight has nobody to supply it.

Reactant lifts every device-resident value reachable from a traced closure into an argument, whether the program reads it or not. A Device field therefore becomes an argument of the exported module that the bundle declares no name for and no client can pass. Measured on a bundle of N weights and one input whose experiment carried two Device arrays and whose layer state carried an RNG seed:

Execution supplied N+1 arguments but compiled program expected N+4

Both readable halves of that bundle were correct. The manifest declared its one input, the safetensors file held its N weights, and the three surplus arguments existed only inside the compiled graph, where nothing was looking.

So the export view FREEZES rather than filters, and that is deliberate. A Device field the exported forward never reads (a loss weight, a class-weight vector, a soft-target matrix) becomes a host value nothing reads, and no constant is emitted for it at all. A Device field forward DOES read bakes into the graph as the value it held at export. Both are the right answer for an artifact that is a fixed function, and neither one adds an argument.

Filtering could not have worked, which is worth stating because dropping the loss-only fields is the obvious design. The framework cannot know which Device fields forward reads: a hook reaches them through ev inside its own body and nothing declares it. Dropping the fields forward does not read would leave every field it does read still lifted, so the defect would survive on precisely the models whose field carries something that matters.

The cost, stated rather than hidden. A frozen Device array is a constant in each compiled module, so a large one is paid for once per batch size in artifact size. A Device field is meant for knobs and per-class statistics, so it is small in practice; an experiment carrying something big enough to matter should override this method and say why.

See also compile_view, which is what a TRAINING trace sees and which leaves Device fields device-resident on purpose.

finalize_metrics(e, acc, split) -> NamedTuple

Host-side, once per split per epoch. Default is identity.

Needed because a derived metric is not the mean of per-batch values:

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

split is the Symbol naming the split, so this is where branching between validation and testing belongs: it is free here and costs a second compiled program in metrics.

finish!(lgr, status) -> nothing

Finalize. status is :completed, :early_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 stamp every failed run completed.

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

Required. One definition of the model's computation, three consumers: training traces forward then loss, validation traces forward then metrics, and predict is forward alone.

Discipline: forward takes inputs, not targets. That is what makes predict possible. A model genuinely needing targets in its forward (teacher forcing) cannot be predicted from inputs alone, which is true rather than a limitation.

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

outputs is passed on as ONE positional argument. Whatever goes in the first slot of this return is exactly what loss, metrics and train_metrics receive in their second; nothing is splatted or unpacked in between. For more than one output, return a NamedTuple and destructure it by name downstream:

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

A Tuple or a nested structure works too, since the framework walks the output tree with Functors, but a NamedTuple is what export_outputs names leaves by.

Every array leaf of the output tree must have the batch dimension last, and must be an array. Both follow from the short final eval batch, which is padded to the compiled width and sliced back: the framework asserts the last axis is the batch rather than slicing the wrong one, and a scalar you already reduced over the batch has nothing to slice. Reduce in loss or metrics instead of returning the scalar here.

State threading is inherent to the contract: for a stateless model st_new is st. The framework owns the train/eval mode switch (Lux.trainmode / Lux.testmode) and users never call either. The st_new returned from an eval-mode call is discarded.

forward is also the exportable program: export uses the same function as predict, so there is no second model definition.

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

Build an experiment from presets(E)[name], with overrides winning over the preset.

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

Every preset key is validated against fieldnames(E), and it is the only correctness argument for presets: today a typo in a splatted NamedTuple is a silent no-op, and a recipe that quietly failed to set the field it names is worse than one that refuses to load.

A preset is field values, so the marker semantics are untouched. GraphConst fields enter the compile cache key as usual, which means switching presets recompiles exactly when the emitted graph really differs: two recipes differing only in Host fields share both compiled programs, and one that changes a class count correctly does not. A preset may set a Device field too; derive wins for anything it computes, which is already true of a hand-set value.

This does not record the name, because it returns a bare experiment and the framework has nowhere to put it. Use Nitro(E, name; ...) for that.

gradient_clip_norm(e) -> Real

The global-norm clip threshold. Default 0f0, meaning off. Also a Nitro keyword defaulting to this accessor, so train!(e; gradient_clip_norm = 1f0) overrides it for one run.

Global norm over the fully accumulated gradient, applied at the top of the optimizer program, never a chain member. A chain member would clip per group, which is a different algorithm producing different updates; a user who genuinely wants that can put an admitted ClipNorm(ω, p; throw = false) in a Level 2 chain.

It is not a Device, not schedulable, and not per group, matching Lightning. It is a trace-time host constant, which is what buys the property that a disabled clip emits no ops at all. The price is that changing it recompiles the optimizer program, and only the optimizer program, so a clip sweep re-pays the cheap compile rather than the expensive one.

0 means off, and a traced threshold could not keep that: with the threshold device-resident, 0 does not disable the clip, it scales the gradient to zero norm, and on an all-zero gradient 0/0 yields NaN silently. "Off" would have to be spelled Inf.

The name matches Lightning's gradient_clip_* family.

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.

A field is Host by default: an unmarked field is recorded here. The GraphConst fields are the complement, setdiff(fieldnames, device_fields, host_fields).

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

The decay coefficient, per group. Defaults are 0 and lambda(e), so by default there is no Decay in the chain at all.

There is one decay coefficient per group and what it decays toward is set separately by that group's decay_anchor. Decaying toward zero and decaying toward the pretrained weights are one rule with different anchors, so the configuration surface is the rule's own two fields rather than two competing coefficients that could contradict each other.

Across groups it can be both at once, and one opt.lambda schedule drives both: the intended setup is :w0 on the pretrained backbone and :zero on the fresh head. That is usually what you want, since both are the same regularization-strength knob, and the binding report names the anchor per group so it is visible.

Decay is decoupled, so the value reaching the rule arrives pre-multiplied by that group's effective learning rate, and an LR schedule modulates regularization strength.

Scheduling opt.lambda while also defining per-group lambda accessors is a setup error naming both, since the schedule supplies the base.

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

The base learning rate, and the per-group bases. Defaults are 1f-3 and learning_rate(e), i.e. ratio 1.0 for every group.

Per-group accessors define ratios; the schedule sets the absolute value of the default group:

η_g(t) = eta_sched(t) * (learning_rate(e, Val(g)) / learning_rate(e))

so η_default(t) == eta_sched(t) exactly, and a backbone at 1f-4 against a default of 1f-3 stays a tenth of it for the whole run. learning_rate(e) == 0 is a setup error, since the ratio is undefined; that is also why the default is 1f-3 rather than 0.

load_checkpoint(ckpt, path) -> CheckpointRecord

Dispatches on the checkpointer, not on the path. A checkpointer with a non-file backend, which is a ten-line implementation, cannot implement a path-dispatched loader at all.

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

Takes a plain Matrix{Int}; the backend adapts it.

THE FRAMEWORK NEVER CALLS THIS. It is for your code. Every other verb in the contract fires from the driver; this one cannot, because a confusion matrix is not something the framework has. It is a user metric, accumulated by the count === nothing rule, and what reaches a phase monitor is the finalized scalars rather than the raw accumulator, so only the experiment knows both that a matrix exists and what its class labels mean.

The shape that works, and the one the first port uses: accumulate it in metrics with a nothing count, stash it from finalize_metrics, and emit it from a register_phase_monitor! on the transition out of EvalStepping, which is exactly when a fresh one exists.

metrics(e, out; y)            = (; confusion = (confusion_matrix(out, y), nothing))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)

The verb is declared here, with the ::Nothing no-op, so that a backend implements it once and every experiment can reach it. That is its whole job.

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

The metric channel. context is "train", "validate", or "data", byte-identical to existing conventions. The framework drops non-finite values before calling.

"data" is the third one and it is deliberately not either of the first two. It carries the data path's per-epoch data_wait_frac, and the two established contexts are each pinned to a cadence that a diagnostic would falsify: a "train" line is exactly one per optimizer step and a "validate" line is exactly one per epoch of metrics. A per-epoch fact about the host data path is a third thing, so it says so instead of diluting either count. A backend that passes the string through, which is all the contract asks, needs no change for it.

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

Free-form single values. The schedule binding report goes through here as log_other!(lgr, "binding_report", str) so it lands in the run's record.

log_params!(lgr, params) -> nothing

Hyperparameters, logged before the first compile, since a compile-time crash would otherwise lose them. Keep the key style flat: dotted keys split comparison tables under one-project-per-model.

log_tags!(lgr, tags) -> nothing

logger(e)

The run's logger. Default _field(e, :logger, JSONLogger()), i.e. an experiment's own logger field if it has one, else the framework's shipped JSON default; also a Nitro keyword. nothing is the documented opt-out and stays the public "no logging" value: pass it as a keyword, or declare a logger = nothing field.

Called once, at the end of setup. That matters because constructing a logger is often a side effect: a file logger opens a handle and a hosted tracker registers a run. One call per Nitro is the contract, so logger(e) = TSVLog(open(...)) is safe. The default JSONLogger is deliberately side-effect-free at construction: it opens its file lazily, after setup has pinned its path to the run's resolved run_dir, the same adoption the default checkpointer gets.

logger_info(lgr) -> NamedTuple

Informational, for humans and for tools, like run_id / run_url: the backend's key identifying parameters, as one plain NamedTuple. A hosted-tracker logger returns its experiment key, URL, workspace, and project; a W&B logger its run id, URL, project, and entity; the shipped JSONLogger its file path. The ReactantNitroKaimonGateExt nitro_logger tool renders exactly this table for a running experiment.

The default is (;), and that is the whole contract: a backend that has nothing to say says nothing, and one that does implements this one small method. The docstring convention is to include the run_id / run_url values under those names when the backend has them, so a tool rendering the table shows the same two canonical identifiers every hosted logger reports.

It must return plain serializable data, a String, Number, Bool, nothing, or arrays/NamedTuples of those, never the live backend object. This is the same rule as logger_state, for the same reason: the table is meant to cross boundaries (a tool reply, a report), and a live handle is useless there. Unlike logger_state this is not stored in the checkpoint record: it is live information about the running experiment, so the record keeps only the two identifiers that existed before this verb.

logger_info(lgr::JSONLogger) -> NamedTuple

The one identifying parameter of a file logger: where the lines land. path is the resolved metrics.jsonl path, or nothing for a pathless default that setup has not adopted yet.

logger_info(nitro) -> NamedTuple

The run's logger's key identifying parameters, read straight off a running experiment: logger_info(nitro.logger). The shipped default JSONLogger reports its path; a hosted backend reports its URL, experiment key, workspace, and whatever else it extends the verb with. The ReactantNitroKaimonGateExt nitro_logger tool renders exactly this table.

logger_state(lgr) -> Any

Machine-readable resumption state, opaque to the framework and backend-specific. Default nothing, meaning "I have no resumable state", in which case nothing is stored and no reattachment is attempted.

Most backends can continue an existing experiment and each needs different information to do it: a hosted experiment key, a W&B run id plus project and entity, an MLFlow run id plus a tracking URI, a directory for a file logger. The framework must not know the shape of that, so it round-trips it opaquely: logger_state on checkpoint, reattach! on resume.

It must return plain serializable data, a String, NamedTuple, or Dict, never the live backend object. The record goes through JLD2 and outlives the process, so a stored handle would be either unserializable or dead on arrival. This is the one contract detail a backend author is likely to get wrong, so the error names it.

The pair is self-describing, which is how it stays loud without burdening simple loggers: if this returns anything but nothing, reattach! is required and a missing method is the usual MethodError. A ten-line file logger defines neither and keeps working.

This is separate from run_id / run_url, deliberately. Those are informational and displayed; this is for machines and never is. An earlier draft used run_id for both jobs, which works only for backends whose entire resumption state happens to be one identifier.

loss(e, outputs; <declared batch fields>) -> scalar

Required. The scalar Enzyme differentiates, traced inside the gradient program together with forward.

outputs is the first element of forward's return; the framework has already stripped st_new. Do not unpack it again: first(outputs) is element one of the prediction array, and every operation after that stays broadcast-legal, so the run would train on one sample out of the batch without raising.

Keyword routing is forward's. loss's own router is also what the default validation metric reuses when an experiment defines no metrics.

manual_training(e) -> Bool

The manual-mode flag. Default: whether the experiment type has a train_step method, which is how defining the hook selects manual mode. Override to false to keep the method in source while using the automatic loop:

ReactantNitro.manual_training(::MyExp) = false

Read once, at construction, and frozen: the driver a Nitro uses is fixed at construction, so flipping this on an existing handle is inert and reported by fixed_config_report rather than applied silently.

max_epochs(e) -> Int

How many epochs to train for. Default _field(e, :max_epochs, 1), i.e. an experiment's own max_epochs field if it has one, else 1.

It is also a Nitro keyword defaulting to this accessor, because raising it on resume is the normal case. Resolution order is keyword, then a user method on the experiment type, then a field on e, then the framework default; the binding report is where a reader sees which source won.

This runs host-side against the real e, never against compile_view's stripped view, which is why the field is normally marked Host: a driver knob has no business in the compile cache key.

A bare experiment therefore finishes after one epoch. Stated because a first run stopping there reads as a bug otherwise.

metrics(e, outputs; <declared batch fields>) -> NamedTuple of (sum, count)

Traced, once per eval batch. A metric reports its own numerator and denominator; the framework adds them up and divides at the end:

metrics(e, outputs; lab) = (; err = (sum_abs_err, n_items),                              acc = (n_correct,   n_images))

A framework-supplied sample count would be the wrong denominator, since different metrics have different natural ones (per-sample, per-image, per-object). count === nothing means accumulate by summation without dividing, which covers confusion matrices.

metrics never sees padding: the framework slices a padded short final batch back to n_real before calling it, so there is no mask batch field and no silent-miscount failure to guard against.

When the user defines no metrics, the framework reports validation loss, as val_loss = (loss(e, outputs; R_LOSS(batch)...), 1), reusing loss's own router. That is framework behavior rather than a default method, and the distinction is load-bearing: see this file's header. The count of 1 makes it the mean over batches rather than over samples, which is a reporting-only difference and never touches training numerics.

The :validation / :testing distinction lives host-side, in finalize_metrics, because as an argument here it would produce two compiled programs even when the user's code ignored it.

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

Where a metric hook runs. hook is :metrics or :train_metrics. Defaults: :host for metrics, :device for train_metrics.

The default follows the cadence, which is the whole argument. train_metrics runs once per micro-batch, so running it on the host means transferring a full output batch every micro-batch, which for anything image-shaped dwarfs the scalars you wanted. metrics runs once per eval batch, once per epoch, against a training epoch that has just done thousands of steps, so the same transfer is noise. Tracing metrics by default would be optimizing the cadence that does not need it, at the cost of the one that does.

And the flexibility runs the other way. A traced metric must be expressible as a traced program, which rules out a great deal that is ordinary in evaluation code: data-dependent control flow, a matching or assignment step, connected components, sorting with tie-breaking, anything reaching a library that knows nothing about Reactant. Those are common in validation and rare in a per-step diagnostic. Computing evaluation metrics on the host is the common case: an evaluation metric is usually ordinary Julia over transferred arrays, and the transfer it costs is noise against the epoch it follows. That is why the host default is the workable one rather than a concession.

Both hooks accept both values, so the choice is the user's:

# A per-step diagnostic too expensive to trace, on a small model where the transfer is affordable.ReactantNitro.metrics_residency(::MyExp, hook) = hook === :train_metrics ? :host : :host# A validation metric that IS traceable, on a large eval set where the transfer is not free.ReactantNitro.metrics_residency(::MyExp, ::Symbol) = :device

What changes for the hook author is what outputs is: device arrays under :device, host arrays under :host. Everything else is identical, including the keyword routing, the (sum, count) contract, and the guarantee that a metric never sees padding.

What changes for the framework is the compile cache. A :host hook is part of no program, so its primary_world is not in the cache key and editing it recompiles nothing. That is the point: it is what makes adding a diagnostic mid-session free.

n_devs(e) -> Int

Local devices to shard the batch over. Also a Nitro keyword. 1 skips the mesh entirely.

The default is every VISIBLE device, length(Reactant.devices()), so a host with four GPUs data-parallelizes without being asked and CUDA_VISIBLE_DEVICES is the supported way to restrict a run. That matches how these hosts are actually driven. With no GPU visible Reactant reports a single CPU device, so a CPU run defaults to 1 and skips the mesh, which is why the test suite is unaffected.

A session-level pin from setup_devices!, which is the nitro_setup tool in a Kaimon session, beats the experiment field and the default. Once a session pins n_devs = 2, every run in that process shards over 2 devices until the pin changes, whatever an experiment declares. An explicit n_devs keyword on Nitro still wins for that one run. Without a pin, the experiment's own field wins over the default.

Note the batch size is GLOBAL and gets split across the mesh, so adding devices buys throughput and does not change the effective batch.

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

Whether this parameter leaf is excluded from weight decay entirely. Default default_no_decay, so a bare experiment excludes biases and norm affines and decays everything else, which is the conventional policy and is exactly what the framework did before this hook existed.

It receives the leaf itself, not merely its keypath, because a useful exclusion rule needs both: the keypath says which parameter, and the array says what shape it is. ndims, size, and eltype are all fair game. This costs nothing, because the hook runs host-side, once, at setup, and never enters a trace, so it can do anything ordinary Julia can.

Four modes, from one hook:

# 1. DEFAULT: write nothing at all.# 2. ADD to the defaults, which is what most experiments want.ReactantNitro.no_decay(::MyExp, ks, x) = default_no_decay(ks, x) || (:fc in ks)# 3. REPLACE them with your own system.ReactantNitro.no_decay(::MyExp, ks, x) = (:fc in ks) && (:weight in ks)# 4. DISABLE exclusion entirely, decaying every parameter.ReactantNitro.no_decay(::MyExp, ks, x) = false

Composition is a plain || against an exported function rather than an implicit merge the framework performs behind you, so what a run actually excludes is readable in the user's own source.

Exclusion dominates the anchor, for free. An excluded leaf gets a mask of 0, and Decay computes no_decay_mask * lambda * (x - anchor), so the term vanishes whether that group's decay_anchor is :zero or :w0. "Excluded from either" needs no separate mechanism.

Exclusion is per LEAF; the coefficient and the anchor are per GROUP. That split is deliberate: biases and norm affines occur inside every group, so expressing their exclusion through param_group would force a group split that also silently splits the learning-rate ratio, which is a different knob. Groups stay the unit at which optimizer behavior is defined; this is the one per-parameter refinement.

Not in the compile cache key, deliberately. The mask is a parameter-sized device buffer passed into Decay as a value, so its contents change no graph and a world entry for this hook could only ever fire spuriously, which is the same reason accum and the clip carry no world entry either. It is resolved once at construction, so revising it takes effect on the next Nitro, like everything else a handle freezes.

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

The fields of rule type R that may not be scheduled, and therefore are not promoted to device residency by to_device_rule. Default ().

One declaration drives both, and it must: a field that cannot be promoted cannot be scheduled, and a field that is scheduled must be promoted. Declaring them separately would let them drift.

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 excluded because it is a Tuple and a scalar schedule cannot produce one; p because it is structural, and promoting it breaks _norm's ::Real dispatch, which is what once made ClipNorm look unable to trace at all. The general rule: any parameter-sized rule field is non-schedulable, since scheduling one would push a parameter-sized buffer to device every step.

Declare one method per concrete rule type, never a Union. A Union method is shadowed by any more specific one, silently and with no ambiguity warning.

The set for a chain is the union over its rules, computed automatically, so a custom chain declares nothing.

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

The two upper levels of the three-level optimizer surface.

Level 0: declare nothing and get RAdam, with per-group hyperparameters from learning_rate and lambda. RAdam rather than Adam deliberately: it rectifies the adaptive variance term over the first steps instead of leaving the user to hand-tune a warmup that does the same job by feel. Adam is one line away at Level 1. The cost is a hard floor of Optimisers 0.4.8, which is the first release whose Reactant extension can trace RAdam.

Level 1, pick the rule. The hook returns a rule type and the framework constructs it, splatting resolved values in by field name, so every configured hyperparameter is applied by construction:

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

The framework-composed Decay tail, per-group accessors, w0 capture, and no-decay mask are untouched. Level 1 rejects any rule declaring a lambda field (AdamW), because that rule plus the framework's own decay tail would decay twice, silently, and opt.lambda would then resolve against two rules in one chain.

Level 2, supply the chain. Same function, higher arity, receiving the group and its already resolved, already device-converted hyperparameters:

optimizer(::MyExp, group::Symbol, hp) = Optimisers.OptimiserChain(    Optimisers.RAdam(; eta = hp.eta, beta = hp.beta),    Decay(hp.lambda, hp.anchor, hp.no_decay_mask),)

hp is a NamedTuple with eta (this group's effective learning rate), lambda (already multiplied by that eta), anchor (a flat device array or nothing), no_decay_mask, and one key per scheduled rule field. group is a bare Symbol here and a Val in the accessors: the accessors dispatch, this branches host-side.

Level 2 owns hyperparameter application, and the framework will not check that the factory read what it was given. A factory building RAdam(; eta = hp.eta, beta = hp.beta) while epsilon is scheduled passes the name check and the schedule then does nothing for the whole run. This is the one place the framework knowingly permits a silent no-op rather than a loud error, and it is confined to the level a user opts into explicitly. It is not invisible: the binding report lists, per Level 2 group, which scheduled values it found in the returned chain and which it did not.

No path accepts a fully constructed chain with baked hyperparameter values, because rules carry tracked device scalars rebuilt every step.

param_group(e, keypath) -> Symbol

The parameter group a leaf belongs to, as a function of its keypath. Defaulted accessor, not mandatory dispatch: the default returns :default for every keypath, which is the single-group case and is what most experiments want.

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

The group is the unit at which optimizer behavior is defined. :default is group 1; remaining groups follow first-appearance order, which is stable for a fixed model and bakes into the compiled program. Per-layer groups (Symbol(join(ks, "_"))) are available if finer resolution is ever wanted, at a cost that is visible since G is a trace-time constant.

parameters(nitro) -> ps

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

phase(nitro) -> Phase

The run's current phase. Phase transitions also reach the monitor registry; this reads the latest one.

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

Inference: forward alone, in eval mode.

It takes a batch NamedTuple or anything iterating them, which is the same data-source contract as the data contract, so a DataLoader works unchanged. It does not take a bare array: forward is keyword-routed from the batch schema, so the framework needs the field names, and the error for a bare array says exactly that, naming the fields forward declares. Only the fields forward declares are required, so a prediction batch does not need labels; that falls out of routing rather than being a special case, and is what lets one loader serve training and inference.

Returns. 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, so predicting over a large set does not materialize every output at once; collect it if you want them all. Host rather than device because the caller is leaving the framework, and a device array that outlives its run is a footgun.

The eval tests pin both halves: a partial batch returns exactly n_real outputs, and 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. The contents belong to the model; the mechanism belongs here. Default is empty, so a type that declares none loses nothing and "an experiment defining only the four required hooks trains with nothing passed" is intact.

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),)

Entries are PARTIAL, and usually are: a recipe states what it changes and everything else falls through to the struct's defaults, exactly as the ad-hoc tables this replaces already do.

Why the framework has an opinion at all, given three models each solved this privately in about ten lines: none of those tables is visible to the checkpoint record or the logger, so a result cannot say which recipe produced it. That, plus validating a key rather than letting a typo be a silent no-op, is what earns the surface. It is not a reproducibility mechanism; provenance already reproduces any run exactly from its git SHA.

Inheritance is merge on NamedTuples and needs nothing from the framework: wide_v2 = merge(wide_v1, (; max_epochs = 80)).

ReactantNitro.progress_counter() -> Int

A monotonic count of units of work this process has completed: one per optimizer step in the training loop, one per batch in the eval loop. Only ever increases, and only its CHANGE is meaningful.

A stall watchdog needs this, and the phase alone cannot give it. A phase deadline measures from the phase transition, and a run stays in TrainStepping for a whole epoch, so a budget meant as "how long may one step take" is silently applied to the entire stretch: an epoch longer than that budget is then killed while perfectly healthy. Comparing this counter across two observations turns the same budget into "time since the last completed unit of work", which is what such a budget is nearly always meant to express.

It covers evaluation as well as training deliberately. Nothing on the handle advances during an eval loop, since the batch index is local to it, so a monitor keying off current_step would leave a long validation or test pass looking motionless.

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 of them. Each entry carries file (a basename, not a path), epoch, score (the validation metric the retention rule ranked that checkpoint by, or nothing for a run with no val split) and stop_reason (nothing until a run records how it ended).

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

This is the cheap question, and it is the one worth asking first. The manifest is a single small file the checkpointer rewrites next to every record, so ranking a run's checkpoints, or finding the newest one to resume from, costs one read instead of deserializing a parameter tree per file in the directory. checkpoint_info is the other half of the pair: go to a record only for the fields the manifest does not carry, and not in a loop over a directory.

Returns an empty vector when dir has no manifest. That is a directory no run has written to yet, which is an answer rather than an error, and it is what lets resume = :auto run in a fresh one.

reattach!(lgr, state) -> nothing

Restore a logger onto its previous experiment, called after the user constructs their logger and before the run starts and before any metric is logged.

Required if and only if logger_state returns non-nothing; there is deliberately no default method, because a logger that claims resumable state and cannot restore it is a bug rather than a configuration.

Type mismatch is caught: the record stores the logger's type name alongside its state, and resuming into a different logger type refuses with both names rather than handing one backend's state to another's logger.

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

Register f as an observer of phase transitions. The module-level form applies to future runs; the nitro form targets the live per-run copy and is effective immediately (an earlier design's registry was module-level only, so a monitor registered during a run never fired in it).

This is the hook a heartbeat or a watchdog is written against. The framework ships neither, and that is the point of publishing the signal: what counts as "too long" is site policy, and a timeout table belongs where the hardware and the operational rules are known rather than in a general framework. What ships here is the transition, on time and in full, so that whatever consumes it can live entirely outside this package.

The signature is f(phase, step, epoch, info) with info::NamedTuple, not keyword arguments. do-block anonymous functions cannot accept keyword arguments at all, and in do syntax a semicolon separates the argument list from the body, so the keyword form could not be written the documented way.

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 Union{Int,Nothing}, because neither is always defined: both are nothing until the first epoch begins, which is what a monitor observing a standalone validate or a transition before the loop sees. info carries at least (; nitro, logger, is_rank0), plus metrics on the transition out of EvalStepping, which is where the finalized numbers first exist, and may grow.

Registry rules: error isolation (a throwing monitor never kills a run; it is caught and warned about once per monitor rather than per event), per-run scope (the module-level registry is copied into the run at train!), and registration order, documented: module-level monitors fire in the order they were registered, then the run's own.

The per-run copy is a snapshot, which is what "applies to future runs" means for the module-level form: unregistering a module-level handle mid-run leaves the live run's copy firing, and takes effect on the next train!. Use the nitro form to reach a run in flight.

The framework fires on transition; sustained liveness is the monitor's job. This is forced by XLA: a compile is a blocking foreign call, so nothing can emit from inside it, and that is exactly the window where a watchdog most needs evidence of life. A monitor that writes only when this fires will look hung during a normal compile, so one that has to prove liveness needs its own task, on the :interactive threadpool, since the :default pool is the one blocked in the foreign call. Compiling is published so that such a monitor can widen its patience rather than guess.

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

Render a few examples from the real pipeline, and return the paths written.

predictions defaults to false, which is the data gate: no forward pass, so no compile. That matters because Nitro(e) compiles nothing and the first forward costs an EvalCompiling phase, which is slow by design. Paying it to draw figures from untrained weights is the one render nobody wants.

You say how many BATCHES, not how many samples. There is no sample cap and so no interaction between a cap and a short final batch. The cost is worth knowing: at batch_size = 64, batches = 1 writes 64 figures. Pass render's batch form a narrower batch if that is too many.

It pulls batches, not epochs, from the head of the split.

Three entry points, one driver, none of which needs train!:

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

Ordering is the loader's, not the framework's: this takes the head of the split in the source's own order, and PrefetchIterator's iterate is a passthrough so nothing here perturbs it. If you want validation renders that diff cleanly across data-prep changes, your validation loader must be deterministic; the framework cannot make it so.

request_stop!(nitro) -> nothing

Ask the run to stop, from a REPL or from a phase monitor. It sets the same flag EarlyStopping sets, checked once per epoch after validation.

Stopping is graceful: the run finishes the epoch, validates, checkpoints, finalizes the logger, and exits through the normal Done path. Aborting mid-epoch would skip exactly the steps that make the run useful.

run_dir(e) -> String

The run's output directory. Default _field(e, :run_dir, joinpath("runs", string(nameof(typeof(e))))); also a Nitro keyword. Checkpoints, the manifest, and resume = :auto all resolve against it.

Note the name is shared with run_dir(nitro), which reads the value a run actually resolved to. They are the same concept at two moments: what the experiment asks for, and what this run got. Dispatch separates them.

run_dir(nitro) -> String

The run's output directory: checkpoints, the manifest, and resume = :auto all resolve against it. It is the one path concept in the framework, and defaults to joinpath("runs", string(nameof(typeof(e)))).

run_id(lgr) -> Any

Informational, for humans and for reports. It stays in the checkpoint record so a checkpoint can be traced back to its experiment without deserializing anything, which is why the ::Nothing method cannot be omitted.

run_url(lgr) -> Any

Informational, as run_id.

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

The checkpointer decides internally whether an epoch warrants a write, so policy and action collapse into one hook.

Writes go through ReactantNitro.with_io_retry, to a temporary path renamed into place atomically, and rotation deletes the displaced file only after the new one is durably in place.

Four arguments, and e is deliberately not among them. The filename comes from the checkpointer's name, bound to the experiment's checkpoint_filename at setup, rather than from an experiment passed in here: the record's host-value discipline exists to keep device-resident leaves away from the function that serializes, and the experiment is the one object in a run that holds them by design.

save_figure(e, fig, stem) -> path

Write what visualize returned, and return the path actually written.

e is in the signature so that writing this method is not type piracy. Without it the one method every model must write is save_figure(::Makie.Figure, stem), which owns neither the function nor the type; two model packages defining it and loaded in one session overwrite each other silently. With e a model package owns MyExp, the specialization is ordinary, and two models can save differently. It is also the only place per-experiment format and resolution can live, and it keeps this hook consistent with every other one in the framework, all of which take e first.

stem carries no extension, because the framework has no opinion about file formats and a framework handing over a path ending in .png has one: that suffix is what every plotting backend dispatches format on.

The return value is what render reports. A backend that writes a video, three files, or a directory of frames says so, rather than having the framework report a single path it guessed at.

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

Every entry is a factory of the horizon. Default (;), i.e. everything 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    epsilon    = 1f-8,                                         # a bare Number is a constant)

The framework calls f(total) once, at setup, then sched(step) once per optimizer step. A bare Number normalizes to _ -> (_ -> value).

The schedule belongs to the experiment by default, because it is part of the recipe. The Nitro keyword of the same name is the per-run override and replaces wholesale; it does not merge. To merge, say so: schedules = merge(schedules(e), (; eta = ...)). The binding report names the source of every entry.

A schedule key is either a Device field name or a rule field name, resolved against the union of the two. opt keys are applied to the optimizer, device keys are written into the experiment and reach traced code as e.field; there is no third destination. Both names are reserved at the top level and are how an ambiguous key is qualified:

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

Rule field names are the rule's actual field names, so the learning-rate key is eta rather than lr.

A nested opt key is a parameter-group path. opt = (; backbone = (; eta = ...)) names the :backbone group and binds that group's chain only, so per-group learning-rate curves work: the backbone anneals while the head warms up. The path value is the BASE curve for those groups and the per-group ratio is still applied, η_g(t) = opt.backbone.eta(t) · (learning_rate(e, Val(:backbone)) / learning_rate(e)), exactly like a bare key; a group with no path key falls back to the bare key, then to its base rate. Paths are ONE level (group.field), because parameter groups are flat, and qualified-only: an unqualified dotted key is a resolution error, since the automatic loop cannot know a top-level key is a group path. Manual mode is deliberately asymmetric: its path values are absolute, because manual mode has no base rate and no ratios.

step is the optimizer step, not the micro-batch. With accum = N they differ by a factor of N, and confusing them shifts the whole curve by N. The horizon is total = max_epochs * div(steps_per_epoch, accum), an exact division because a training batch count not divisible by accum is a setup error.

Schedules are host-side, so ordinary Julia control flow is fine inside them. A schedule need not be a pure function of the step, but one closing over training state does not resume exactly, since resume restores the step counter and not the closure.

ParameterSchedulers.jl is a documented recommendation with no dependency, not even a weakdep, because the framework never dispatches on a schedule.

seed(e) -> Int

The rng seed. Default _field(e, :seed, 42); also a Nitro keyword.

A seed field must be Host, and unmarked already is: the seed is kept out of the config hash so that a seed sweep shares one compiled program, and a GraphConst field of that name would enter the compile cache key and recompile per seed. That is checked at setup rather than left to the reader, since the whole point of the sweep is that it is cheap.

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

Write a new value into a Device field of a live handle. Provably does not recompile, which is the whole point: it is how a device sweep or a changed inference threshold reuses the programs already in the compile cache instead of rebuilding a Nitro.

The guarantee is structural rather than hopeful. The cache key skips device_fields when hashing, and graphconst_field_hash seeds on Base.typename(T) rather than hash(T) precisely so that re-parameterizing an experiment for device residency does not move the hash, so a new device value in a Device slot leaves every key component untouched. This function asserts that before it commits, and the assertion is not decoration: typeof(compile_view(e)) is args[1]'s type at every trace site, so a value whose device type differs by even an element type would move the key and buy a recompile silently.

Two things are refused, both because they would break that guarantee rather than out of caution:

• A field that is not Device. A GraphConst field bakes as a trace-time constant and is hashed, so changing it genuinely is a different program; a Host field reaches no trace at all. The error names which case it is and points at the rebuild, which is cheap because the cache is module-level and a fresh Nitro hits every entry the change does not invalidate.
• A different type or size. Element type would move typeof(ev) and therefore the key. Size would not, since _shape of a struct is nothing, which is worse: the key would still match and the shape mismatch would surface inside XLA at call time, exactly the gap the key lists shapes for.

A scheduled field is also refused, because the per-optimizer-step rebuild rewrites the scheduled entries from the schedule, so a value written here would survive until the next optimizer step and no longer. Silently losing a write one step later is worse than refusing it.

n = Nitro(e; checkpoint = "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 compile, any iterationend

See also device_value to read one back, and Device for what the marker means.

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

Configure this process's accelerator for every subsequent run and report what is now in effect, as (; backend, visible, n_devs, pinned). The Kaimon tool nitro_setup is exactly this function, so a REPL session and a Kaimon-hosted session configure identically.

Calling it is optional. A session that never calls it runs on Reactant's default backend with n_devs = every visible device, length(Reactant.devices()), which is what a CPU machine (one visible device) gets automatically. This function exists to be explicit and to fail fast.

backend passes through to Reactant.set_default_backend and names a Reactant backend: "cpu", "gpu" (whichever of CUDA/ROCm is available), "cuda", "rocm", "tpu". Omit it to leave Reactant's default: the highest-priority working backend, chosen once per process at the first device access (GPU where one is visible, else CPU).

n_devs pins the device count for this process, validated eagerly against the VISIBLE devices through the same check a Nitro construction runs, one session earlier. The pin replaces the default and beats an experiment's declared n_devs field; an explicit n_devs keyword on Nitro/train! still wins for that one run.

One process, one XLA. A Julia process initializes its XLA/PJRT client once; n_devs never creates clients or processes. It slices the already-visible device set to build the mesh, so to restrict which GPUs a session sees, set CUDA_VISIBLE_DEVICES before starting the process: visibility is fixed at the first client access. Asking for more devices than are visible is an error, and the error names CUDA_VISIBLE_DEVICES.

The batch size is GLOBAL and gets split across the mesh, so adding devices buys throughput and does not change the effective batch.

Calling with no arguments changes nothing and reports the current configuration: the backend in use, how many devices are visible, and the n_devs the next run will use.

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 the user owns, one per network or group:

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))

Called once at construction, in place of the automatic loop's own optimizer-state build, after device conversion and build_model. The framework normalizes the result to device residency through the to_device_leaf walk and asserts the no-host-Number property on it, so a rule's scalars are traced and an integer step counter that freezes under trace is promoted before it can. The result is what train_step receives as opt_state.

Rules are fixed for the run: there is no per-step rebuild, so rule-field schedules are a setup error. Fixed rules are safe across steps because the closure's program does not donate their scalars (measured).

should_stop(es, epoch, metrics) -> Bool

The stopping policy. Custom policies implement this for their own type; a ::Nothing method returns false, which is how "no early stopping" stays the default with no branch in the driver.

patience counts epochs without improvement and min_delta is absolute, both matching Keras and Lightning: an epoch improves when it beats the best seen by more than min_delta, and anything else, including an equal or slightly better value, counts against patience.

The metric is read through check_control_readback, because this is control flow and every scalar the framework branches on must be validated: a failed BufferToHost on this stack returns garbage without raising, and a garbage value here either truncates a healthy run or lets a stalled one continue.

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

The SITE's half of a bundle's provenance: repository state, collected at root. The second verb a backend answers, and the reason it is a verb at all is the one export_provenance gives: a git commit, a tree hash and a working-tree patch are site policy rather than framework knowledge, and a framework that shelled out to git would be asserting that the process's working directory is the model's repository.

The caller names the root, so nothing is guessed. That is what makes this compatible with export_provenance's refusal to look: the framework still does not decide what repository a model lives in, it forwards a root it was given to a backend that knows how to read one.

It dispatches on the backend rather than being hardcoded because the artifact format owns what it can record. ReactantServerBundle answers with ReactantServerExport.collect_provenance, whose git_diff the writer materializes as working_tree.patch in the bundle. On a dirty tree that patch is the only thing tying the artifact to the code that produced it, and it is a multi-line unified diff, which is exactly the shape that cannot ride a flat name=value list. That is why this is reached through a ROOT rather than through the provenance dictionary.

There is no default method, deliberately. A backend that cannot collect site provenance must say so rather than return an empty dictionary, because a silently empty result is the failure this whole surface is designed against: a bundle that looks complete and is untraceable.

states(nitro) -> st

The layer state, including running statistics. Plural because state reads as the handle's own state, which is the whole Nitro.

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

One tree-level optimizer step for use inside a train_step closure: Optimisers.update under trace, returning the new parameter tree and the new optimizer states.

The rules are stripped from the return, and the driver re-attaches the rules it handed in. A measured rule of sharding: a replicated scalar, which is what a rule's device hyperparameters are, cannot leave a compiled program as an output on a mesh, though arrays can, and that is the same rule that makes the automatic opt_program return states. opt_state must be device-normalized, which setup_optimizers does; call this once per network or group, passing that subtree:

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

The subtraction mirrors apply_group in preserving eltype(x), so the parameter tree is a type fixed point across the step, which is what keeps the closure's program on one cache entry.

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

Train. train! blocks and returns the Nitro.

train!(nitro) takes no keywords; every keyword belongs to the Nitro constructor, and train!(e; kwargs...) is pure sugar for train!(Nitro(e; kwargs...)). Three ways to reach the handle: the return value, run_ref::Ref{Nitro} filled before the loop starts, or info.nitro inside a phase monitor.

Re-training one Nitro under a different stopping rule means constructing another one, which is cheap: the compile cache is module-level, so a second Nitro over the same experiment reuses every compiled program.

On divergence: there is no non-finite rollback. A non-finite loss stops the run with an error naming step and epoch, at whatever cadence the loss is already read. Recovery is resume from the last checkpoint with a lower learning rate. The read itself needs validating because this stack has a documented bug where a failed BufferToHost readback returns garbage without raising, so the framework validates every scalar readback it uses for control flow (the loss, the checkpoint metric, the early-stopping metric) and lets purely-logged metrics through unvalidated. That asymmetry is deliberate.

Ctrl+C is a graceful stop. The loop runs on a worker thread when one is available, and ^C interrupts the parked caller rather than the run; this entry point turns it into request_stop!: the step loop breaks at its next boundary, the epoch's validation and checkpoint still run, and the call returns the finished Nitro with stop_reason = :requested, exactly the nitro_stop wind-down. A run that fails surfaces its own exception, not a task wrapper.

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

Traced, once per micro-batch, inside the gradient program. Default (;), which must fold away entirely.

Train and validation metrics have deliberately different contracts: these are scalars per step with no reduction, for an instantaneous diagnostic, while metrics is (sum, count) per batch, accumulated over a set once per epoch. Smoothing is host-side in the user's logger.

Readback is lazy, so a value returned from a compiled thunk is a device array and costs D2H only when read: compute these in-graph every step and read them only on logged steps, which avoids separate programs for logged and unlogged steps. See check_control_readback for the readback hazard that creates.

The primal is exposed to metrics code inside the traced program, not transferred: the step returns (loss, st_new, stats) and only scalars cross the boundary.

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. The automatic loop has no train_step: its step is framework-owned, grad_program plus opt_program plus rebuild_rules sequenced by the driver, and the user never writes one. So the act of writing this method is what selects manual mode; manual_training(e) = false declines it while keeping the method. An experiment that defines neither trains automatically, and an experiment that defines it but is only evaluated never consults it.

Manual mode hands you the whole optimizer step. The framework keeps everything outside the step: the batch stream and prefetch, per-epoch accounting, validation, checkpointing, logging, phases, and early stopping. Inside one call you own the forwards, the backwards, the optimizer steps, and the device metrics.

The call. Once per optimizer step, in train mode (the framework owns the Lux.trainmode switch), on one transferred device batch. The batch's fields are routed from the keywords this method declares, so a field only the closure reads (the GAN's noise, below) is transferred only when declared. The method is traced and compiled by the framework, so it must be a stable named method, never an anonymous closure built per step.

The objective discipline. The gradients inside are computed with backward(f, ps_sub, consts...), whose objective must take every traced value it reads as an argument. Closure captures of traced values are silently treated as constants, which is how a zero gradient happens without an error (measured). A plain struct like the model, which holds no arrays, is safe to capture; parameters, batches, and state go in consts....

A very simple GAN. Generator and discriminator as two Lux chains in one parameter tree, each with its own optimizer; noise z rides in the batch. The generator's backward re-runs its own forward from the Duplicated subtree, and the discriminator's parameters are a Const argument, so the generator's gradient flows through the discriminator's forward as a function of the fake data and never into the discriminator's parameters (the non-saturating formulation):

using Statistics: mean   # or `using Statistics` at the top of the file@experiment struct ToyGAN    max_epochs::Host{Int} = 50endReactantNitro.build_data(e::ToyGAN, dist) = (;    train = [(; x = randn(Float32, 4, 32), z = randn(Float32, 2, 32)) for _ in 1:10])ReactantNitro.build_model(e::ToyGAN, rng) = begin    model = (; gen  = Lux.Chain(Lux.Dense(2 => 16, tanh), Lux.Dense(16 => 4)),               disc = Lux.Chain(Lux.Dense(4 => 16, tanh), Lux.Dense(16 => 1)))    (model, Lux.setup(rng, model)...)endReactantNitro.setup_optimizers(e::ToyGAN, model, ps, st, mesh) = (;    gen  = Optimisers.setup(Optimisers.Adam(2.0f-4), ps.gen),    disc = Optimisers.setup(Optimisers.Adam(2.0f-4), ps.disc))function ReactantNitro.train_step(e::ToyGAN, model, ps, opt_state, st; x, z)    # forwards: the generator's output is the discriminator's fake input. Each sub-network    # applies with its own state subtree; `st` stays the WHOLE tree for the backwards and    # the return.    fake, st_gen = Lux.apply(model.gen, z, ps.gen, st.gen)    d_real, _ = Lux.apply(model.disc, x, ps.disc, st.disc)    d_fake, _ = Lux.apply(model.disc, fake, ps.disc, st.disc)    # discriminator backward: real -> 1, fake -> 0    l_d, g_d = ReactantNitro.backward(ps.disc, x, fake, st) do ps_d, xc, fc, stc        d1, _ = Lux.apply(model.disc, xc, ps_d, stc.disc)        d2, _ = Lux.apply(model.disc, fc, ps_d, stc.disc)        mean(abs2, d1 .- 1.0f0) + mean(abs2, d2)    end    # generator backward: fake -> 1, through the discriminator's forward, never into its params    l_g, g_g = ReactantNitro.backward(ps.gen, ps.disc, z, st) do ps_g, ps_d, zc, stc        f2, _ = Lux.apply(model.gen, zc, ps_g, stc.gen)        d3, _ = Lux.apply(model.disc, f2, ps_d, stc.disc)        mean(abs2, d3 .- 1.0f0)    end    # one optimizer step per network; states out, rules re-attached by the driver    s_g, ps_g = ReactantNitro.step_optimizer(opt_state.gen, ps.gen, g_g)    s_d, ps_d = ReactantNitro.step_optimizer(opt_state.disc, ps.disc, g_d)    return (; loss = l_d + l_g,              ps = (; gen = ps_g, disc = ps_d),              st = (; gen = st_gen, disc = st.disc),              opt_state = (; gen = s_g, disc = s_d),              stats = (; l_d, l_g))endtrain!(Nitro(ToyGAN()))

The return. A checked NamedTuple:

The driver checks the keys, validates loss against the no-non-finite rule, logs (; loss, stats...), and stores ps, st, and the merged opt_state for the next call.

Schedules. schedules works in manual mode with one difference: opt keys are path-bound into the opt_state setup_optimizers returned, because the framework has no parameter groups to bind them to. A bare key (opt.eta) binds every rule with that field, at the absolute value; a nested key (opt.gen.eta) is a path into the tree and binds the rules at that subtree. The driver rebuilds the named rules between calls, type-preserving, so the same program serves every step. device keys work exactly as in the automatic loop. accum > 1 remains a setup error (accumulation is the automatic loop's mechanism; the closure owns its own multi-batch work). The train_metrics hook is unused in manual mode: stats come from the closure itself.

What the automatic loop still does for you in manual mode, namely validation, checkpointing, early stopping, request_stop!, phases, and the device schedules, is unchanged; loss becomes optional (the closure computes its own losses; define metrics for validation, or loss for the val_loss substitution when metrics is also absent).

unregister_phase_monitor!(handle) -> nothing

Remove a monitor registered by register_phase_monitor!, by the handle it returned. The handle carries the registry it was added to, so one verb removes from either without the caller having to say which, and removing an already-removed monitor is a no-op rather than an error.

validate(nitro) -> NamedTuple

Run the :val split and return the finalized metrics. This is what the training loop calls each epoch, and it works on a Nitro that has never trained.

Eval mode throughout: the framework calls Lux.testmode(st) and the st_new returned is discarded, so nothing accumulates during validation. The loop frees each batch's device buffers explicitly after its metrics are accumulated, rather than leaving it to the GC, which is the documented cause of device OOM during validation on this stack.

Like every entry point, this runs its loop on a worker thread when one is available, and Ctrl+C stops the split at its next batch boundary and surfaces the interrupt.

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

Draw one sample. Deliberately shaped like metrics: positional experiment, positional outputs, batch fields by keyword, routed to exactly what the method declares. A reader who knows metrics knows this.

It is called once per SAMPLE, with the batch dimension already dropped. That is the whole boilerplate reduction: you write a function of one example and never write [:, :, :, i] nor reason about batch layout. A rank-1 field yields its element, so a Vector{String} of case identifiers arrives as a String rather than as a zero-dimensional view, which would interpolate into a title as fill("case_B").

outputs is nothing in data mode, and nothing is dispatchable. One generic method therefore covers both jobs, and two methods split them when the figures have little in common:

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

The two methods may declare different batch fields, which is usually wanted since the data figure needs less than the prediction figure. Shared axes are a shared plain function both call.

There is no default method. Visualization being optional means "you need not call render", never "render may quietly do nothing", so a missing method is an error naming the experiment type and which of the two modes was missing.

The return value is whatever save_figure knows how to write. The framework does not look at it.

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

The one verb a backend implements. Everything above it is the framework's and is already resolved by the time this is called; everything below it is the artifact format's and is none of the framework's business.

model is a callable the framework built, with the model(inputs, ps, st) -> (outputs, st) shape a Lux-style tracer expects. It already carries compile_view(e), the resolved routing, the traced preprocess, and the output selection, so a backend traces it exactly as it would trace any model and needs to know nothing about experiments.

Keywords, all supplied by export_model:

• dir, name: where the artifact goes and what it is called.
• input_names, output_names: Vector{String}, in the order the program takes and returns them.
• output_select: maps the raw forward return to the ordered tuple of arrays that ship.
• client_inputs, client_outputs: nothing, or the Vector{ExportSpec} the client side uses.
• postprocess: nothing, or the model.jl source to write into the artifact.
• batch_sizes: each one is a separately compiled program.
• provenance: Dict{String,Any}, already merged (framework first, caller's on top).

A backend is expected to return the path it wrote.

@experiment struct MyExp ... end

One declaration point, generating the struct, a @kwdef-style keyword constructor, the device_fields and host_fields traits, the config_metadata table, a compile_view method, two 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

The three markers are the three field categories the rest of the framework turns on. A Device field is a traced input. A GraphConst field bakes as a trace-time constant and enters the compile cache key. A Host field is invisible to the tracer entirely, and it is the default: an unmarked field is Host, which is the right category for the majority of real fields (driver knobs, dataset-sized state).

Why a macro is necessary rather than merely nice. Each Device field's storage varies independently (scalar to ConcretePJRTNumber, array to ConcretePJRTArray), so N device fields need N type parameters; and after setup a field holds a ConcretePJRTNumber rather than a Device, so the framework cannot recover which fields were marked by inspecting types at runtime. The declaration-time knowledge has to be recorded as a trait.

Generated type parameters. Every Device and Host field gets its own free type parameter, so a Device can hold a host value before setup and a device value after, and a Host can hold its real value in e and a StrippedHost in compile_view(e). GraphConst fields keep their declared concrete type, and the generated keyword constructor converts to it, so n_layers::GraphConst{Int} means what it says whether or not the experiment happens to declare a Device alongside it. Parameterized fields are not converted, since the point of the parameter is that what the field holds changes.

Docstrings are written as bare strings above each field. Comments are stripped by the parser and are never visible to a macro, so # cannot carry a description. Per-field docs land in config_metadata and in the generated type docstring.

Display shows values and never a device buffer. The generated show prints one line per field with its marker and its value, summarizing an array as its eltype and shape and recursing through tuples and NamedTuples to do it. That matters because Device{T} takes any T, and a read-only buffer for an hlo_call lives in a Device{NamedTuple} or Device{Tuple} of weights: under Julia's default struct show, printing such a config prints the arrays element by element. The generated output grows with the field count and never with the model.

The macro is optional. A user may hand-write the struct and define device_fields, host_fields, and config_metadata themselves; those three are exported for exactly that reason. Such a struct keeps Julia's default show, and opts in with the one line the macro expands to:

Base.show(io::IO, e::MyExp) = ReactantNitro._show_experiment(io, e)Base.show(io::IO, ::MIME"text/plain", e::MyExp) = ReactantNitro._show_experiment(io, e; long = true)

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

The counter behind the acceptance 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.

StrippedHost{name}

The value compile_view substitutes for a Host field, carrying the field name as a type parameter so an error message can be built without the sentinel holding any data.

Using one raises, and the error says what to do about it. The sentinel used to define no methods at all and rely on Base's fallbacks, which was loud but useless: you got MethodError: no method matching *(::StrippedHost{:sz}, ::Float32), which names the field only by accident of the type parameter and offers no fix. Since Host became the default, reading an unmarked field from traced code is the common mistake rather than an exotic one, so the high-traffic operations now carry real methods and _stripped_error writes the message.

Stated honestly, and unchanged by that: the sentinel is still accepted by any ::Any signature, still compares with === and ==, still hashes, still survives in a returned NamedTuple, and still returns false from isnothing without raising. Methods improve the message on every path that was already loud; they do not widen detection. A Host value that is merely stored, compared, or passed through still escapes, exactly as before.

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

Fire the monitor registry for phase without recording it on the handle. The one case set_phase!'s "record and publish are one operation" rule does not cover.

It exists for Repl, and the reason is a distinction worth naming: Repl is a property of the PROCESS, not of the run. Every other phase answers "what is this run doing", and phase reads it back, which is why phase(nitro) isa Done after a successful train! and isa Failed after one that raised. That is the handle's record of the outcome and callers depend on it. Repl answers a different question, "does anyone have work in flight in this process", so writing it into nitro.phase would erase how the run ended in order to say something that was never about the run. A monitor still needs the event, because a heartbeat cannot otherwise tell a finished run from a process wedged in teardown.

So the two are split: the run's phase is recorded and published, and Repl is published only. phase(nitro) keeps naming the outcome; the monitor stream carries both.

No transition guard, unlike set_phase!: there is no stored phase to compare against. The depth counter is what keeps this to one event per outermost call.

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

Publish through the MODULE-LEVEL monitors, for the window in which there is no handle to publish through: Nitro(e) itself. build_data runs there, and on this stack that can mean starting and compiling a data server, so the construction is minutes of real work that no run has declared yet. Without this a supervisor is still being told whatever was declared before the call, which under a session that budgets idle time is a budget already counting down.

A monitor adopted into a run later (adopt_monitors!) is the same object, so a Starting published here and a Compiling published through the handle afterwards reach the same observer in order.

ReactantNitro.work_in_flight() -> Bool

Whether any public entry point is currently executing in this process. false means the caller has control, which is what Repl announces.

A monitor needs this as well as the transition, because a transition can be missed: an entry point that throws before it has a Nitro to publish through, most obviously a failing Nitro(e), leaves the last published phase in place. A monitor that re-stamps on an interval can reconcile against this predicate and correct itself, which is the difference between "briefly wrong" and "wrong until the process exits".

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

Move the run to phase and publish the transition. Setting the field and firing the registry are one operation on purpose: a phase recorded but not published is a monitor that misses it, and this is the only writer.

Fires on transition only. The phase leaves are singletons, so re-entering the phase you are already in is a no-op rather than an event per batch.

Any extra keywords are merged into info, which is how the transition out of EvalStepping carries metrics.

ReactantNitro.graphconst_field_hash(ev) -> UInt

A hash of compile_view(e)'s GraphConst fields only, that is setdiff(fieldnames(typeof(e)), device_fields(typeof(e)), host_fields(typeof(e))).

Both exclusions are needed and neither subsumes the other. compile_view strips Host and leaves Device in place, so hashing the view itself would hash device scalars that change every step and every optimizer step would miss the cache and recompile.

Hashing rule: hash each included field. This hashes the instance of T, never the GraphConst{T} marker, which is a zero-field type the @experiment macro strips: the declared field type is T and the stored value is a T.

That is necessary and not sufficient. Hashing the instance is content-based only if T has a hash method. A config struct holding a Vector has none, so Base.hash falls through to hash(objectid(x), h) and reaches the vector by identity rather than descending into it; the vector's own content-based hash, one level down, is never called. assert_graphconst_hashable refuses that at setup rather than letting it reach a compile, and carries the full argument for refusing over hashing structurally.

ReactantNitro.assert_graphconst_hashable(e) -> nothing

Refuse at setup a GraphConst value that does not hash and compare by CONTENT, because every guarantee the compile cache and the resume check make about configuration rests on hash and isequal answering "is this the same configuration" rather than "is this the same object".

The check is one line of semantics: a value and its deepcopy must hash equal and be isequal. A deepcopy is structurally identical and a different object, so the two agree exactly when the answer comes from contents.

What this catches, and what it deliberately does not. A GraphConst{Vector{Int}} passes: hash(::AbstractArray) is content-based, and using a vector in a configuration is an ordinary thing to do. What fails is a GraphConst whose value is a struct containing a mutable field, because Base.hash has no method for that struct and falls through to hash(objectid(x), h), and objectid reaches a Vector field by identity rather than descending into it. The perfectly good hash of the vector one level down is never called.

Three symptoms follow, and the third is why this is an error rather than a warning:

• A full recompile per Nitro handle. Two identically constructed experiments produce different keys, and the gradient program is measured in hundreds of seconds.
• A spurious resume refusal. check_config_compatible reports a field as changed and prints a diff whose two sides are identical.
• Silent reuse of the wrong program. Mutating the vector in place does not move the key, so the framework serves a program compiled for the old value and the resume check passes. The cache's contract is that the framework refuses to silently reuse a program a change it can see would invalidate; this is a change it cannot see, and it is a second hole beside the documented one.

Why refusing beats fixing it in the framework. A structural hash that descends into any struct without its own hash method would work, and was rejected: it would change the key of every existing experiment carrying a struct-valued GraphConst, so every checkpoint for those models would refuse to resume once, and it would silently paper over a type whose author never decided what equality means. Refusing costs three lines in the model and makes the contract visible.

ReactantNitro.check_config_compatible(record, e, path) -> nothing

The config half of the resume check: identical config continues, changed config errors, and neither silently does the wrong thing.

The comparison is over GraphConst fields only, which is the same set the compile cache is keyed on, and that is not a coincidence: a GraphConst field bakes as a trace-time constant, so a changed one means the resumed run would train a different compiled program than the one being resumed.

Device fields are excluded because they are traced inputs that cannot affect the graph, and they are recorded separately in devices so a change stays visible after the fact. Host fields are excluded because raising max_epochs on resume is the normal case. A derived GraphConst field is in the comparison and a derived Device is not, which is exactly the three-case rule above and needs no special handling here: derive merges into the struct before this runs, so each derived value is already whichever kind it was declared as.

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

The message a user gets when traced code reads an unmarked field. The framework cannot know which fix was meant, so it offers both with their consequences rather than guessing: a value that should bake into the graph wants GraphConst and a recompile per distinct value, and a value that should cross as a traced input wants Device and no recompile at all. Naming only one of them would be advice half the time.

The third option is in there because it is the commonest real answer: most fields are read host-side and should stay unmarked.

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

The validated readback, for the scalars the framework branches on: the early-stopping metric here, and the checkpoint selection metric. This stack has a documented bug where a failed BufferToHost returns garbage without raising, so a value that reaches a decision is checked and a value that is only logged is not. That asymmetry is deliberate: dropping a logged NaN costs a missing point on a chart, while acting on one truncates a healthy run or lets a stalled one continue.

ReactantNitro.with_io_retry(f; attempts, backoff) -> f()

Retry a filesystem operation that failed transiently. This exists because a checkpoint write is the one I/O path in the framework whose failure loses hours of compute.

Checkpoint writes go through here, write to a temporary path and rename into place atomically, and top-K rotation deletes the displaced file only after the new one is durably in place. A rotation that deletes first and then fails to write leaves a run with fewer checkpoints than its retention policy promises.

ReactantNitro.PrefetchStream

The fan-out, as one object so close_stream! has something to dispatch on and one place to stop every task. Iterating it yields the same (host, device) pairs the other two paths yield, which is what keeps the training loop's body identical across all three.

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 is host_batches made literal. The coordinator takes a token before handing out a job and the transfer stage returns one after emitting a batch, so at most that many host batches exist at once across the workers, hostch, and the reorder buffer together. It is a Channel rather than a Semaphore for one reason: teardown closes it, and a coordinator blocked on take! of a closed channel unwinds into its own finally, whereas one blocked on a semaphore would be a leaked task with no way to reach it.

marks is the exactly-once ledger: one byte per batch index, written by the worker that took that job. check_prefetch_delivery asserts every byte is set at the end of an epoch that ran to completion. It is Vector{UInt8} rather than a BitVector because distinct bytes are independent memory and distinct bits of one word are not, and N workers write concurrently.

ReactantNitro.auto_prefetch(collection) -> collection

Wrap every split in a PrefetchIterator at the framework's defaults, unless it is already a PrefetchIterator or a NoPrefetch.

This is the change that addresses the footgun rather than only the symptom. Before it, prefetch was opt-in through a wrapper the user had to remember in build_data, one model's comment asserted the framework was handling it, and for weeks nothing contradicted that. The default now has to be declined rather than requested, and the resolved settings appear in the binding report.

Every split, including the eval ones, which was not always true here: while run_eval iterated its split directly, a worker count on an eval split would have been a number in the report that nothing used. eval_stream is what changed that, and the reason it exists is that evaluation is the phase MORE likely to be starved, not less: a validation step is forward-only, so device time per batch falls sharply while host time per batch does not move.

Called at setup after derive, the schema probe, and the contract checks, so all three see exactly what build_data returned.

ReactantNitro.batch_stream(split, routing)

The training loop's batch source, yielding (host_batch, device_batch) pairs. One loop body serves both paths, which is why the pair is the element type: the loop needs the host batch for the schema and shape checks and for a :host-residency train_metrics, and the device batch for the gradient program.

Without a PrefetchIterator this is a lazy generator and the transfer happens inline, exactly as it did before prefetch existed. With one it is a Channel of device_batches fed by a spawned producer, so the transfer overlaps the previous step's compute.

A producer that throws surfaces at the consumer. Channel's task form closes the channel with the exception, and iteration on the consumer side rethrows it, so a failing loader stops the run with its own error rather than hanging the loop waiting on a channel nobody will feed. That is Julia's own behavior and is asserted rather than reimplemented.

The checks stay on the consumer, against the host half of the pair, rather than moving into the producer. Ordering is then unchanged from the inline path, and an error is raised on the task the user's stack trace is about. The cost is that a batch with a bad schema is transferred before the check fires, which is one wasted transfer on a path that is about to raise.

The three paths, and which one a split takes

1. Fan-out (PrefetchStream), when workers > 1 and the source is fanout_capable. N workers build host batches, one task transfers, the consumer sees the same pair type as always.
2. One producer, a Channel of device_batches fed by a single task iterating the source. What this function used to do unconditionally, and the fallback when the source does not implement the index-addressable trait.
3. Inline, a lazy generator, when prefetch_device_batches == 0: a NoPrefetch split.

prepare is what makes one pipeline serve both phases

The producer's job is "host batch in, device payload out", and the only difference between training and evaluation is what that payload is: training transfers the batch as it stands, evaluation pads a short final batch to the compiled width first. Passing the step as a closure keeps ONE fan-out, one credit window, one joiner, and one teardown, rather than a second copy of the machinery that the tests would have to cover twice and that would drift the first time either is fixed.

The three-argument form is the training one and is what every existing caller writes; eval_stream supplies the evaluation closure.

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

batch_stream for an evaluation pass: the producer pads a short final batch up to batch_size and transfers the padded one, so the pad and the H2D copy both overlap the previous batch's forward instead of running between them.

n_real is deliberately NOT carried through the stream. It is batch_size_of(batch, xfer), a pure function of the unpadded host batch that the consumer already receives, so recomputing it there costs one size read and cannot disagree with what the producer padded to. Carrying it would widen the pair every stage of the pipeline passes around, for one consumer.

Why evaluation wants this at all, given it is the cheaper phase: it is the cheaper phase on the DEVICE. A validation step is forward-only, so device time per batch falls sharply while host time per batch does not move at all, which makes the host:device ratio worse than training's, not better.

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

Assert that batch_at agrees with the source's own iterate, batch for batch. This is the only thing that catches a batch_at whose index units are wrong, and it needs no accelerator, so it belongs in a model package's test suite or in a CPU gate.

One epoch's plan is used for both halves: the sequential pass is taken first (which calls begin_epoch! through iterate) and the indexed pass follows with no re-plan between them.

values = false compares field names and per-field sizes, which is what a loader with stochastic augmentation permits. values = true compares contents and is the strong form; run it on a deterministic split, which for these models means validation.

ReactantNitro.close_stream!(stream) -> nothing

The cleanup half, and the reason batch_stream's consumer runs inside a try/finally: early exit must stop the producer and free the device buffers it is holding, or it leaks a task and device_batches x batch of device memory. An early exit is not exotic here: request_stop! from a monitor, a non-finite loss, and any error mid-epoch all leave the loop with a full channel behind it.

close is what stops the producer: a blocked put! on a closed channel raises, which ends the task. The buffers already in the channel are then freed explicitly, since they are the ones nothing else has ever held.

A generator has neither, so this is a no-op on the inline path.

On the fan-out it must stop several tasks and it must not block. Closing each channel is what unwinds each stage: the workers leave for i in jobs, a worker blocked in put! on the host channel raises, and the transfer task blocked in put! on the device channel raises. It deliberately does not wait on the tasks, because a worker in the middle of a network round trip would hold teardown for seconds; it will raise on its next put! and exit, and nothing it holds is device memory. That is the third reason the transfer lives on its own task: only one task in the pipeline ever holds a device buffer, and it is the easiest one to stop.

ReactantNitro.fanout_capable(source) -> Bool

Whether source implements both halves of the index-addressable trait. Both, deliberately: see begin_epoch!.

Either batch_at shape satisfies the second half. There is deliberately no generic three-argument forwarding method, because one would make hasmethod answer true for every source alive and this check would stop meaning anything; the framework picks the shape at its one call site instead, on whether the plan came back nothing.

ReactantNitro.strip_rules(opt_state) -> states

Replace every Optimisers.Leaf reachable from a user's optimizer state tree with its STATE, preserving structure. The mirror of merge_rules, and the half of the "rules in, states out" contract that runs inside the traced closure: a Leaf cannot leave a compiled program on a mesh because its rule carries replicated scalars, so the closure returns states and the driver re-attaches the rules.

ReactantNitro.merge_rules(combined, states) -> opt_state

The mirror of strip_rules, run by the driver between closure calls: re-attach the rules of the opt_state it handed in to the states the closure returned, rebuilding each Optimisers.Leaf as Leaf(input.rule, returned_state, input.frozen). The two trees must have the same structure; a mismatch is a framework bug rather than a user error, and surfaces as a method error here, at the merge, rather than inside a trace.