ReactantNitro.jl

A training framework for Reactant-first machine learning

Declare an experiment. The framework owns the compiles, the device transfers, the optimizer, the schedules, and the loop.

Tutorial API Reference View on GitHub

⚡

Reactant-first

XLA is the accelerator under the hood. Every compiled program is Reactant, and the framework exists to make Reactant workflows fast, correct, and boring.

🚀

Compile once, train forever

A module-level compile cache hashes the GraphConst fields on your experiment, so only a change that really moves the graph recompiles. Device values never do.

🔁

REPL-native

Revise, rebuild, repeat: the workflow is designed to flow naturally in a REPL session, perfect for testing new ideas on a whim.

🧩

Lux.jl models

Your model is an ordinary Lux model. ReactantNitro builds on Lux the way Lightning builds on PyTorch.

🎛️

Parameters and optimization

Parameter groups, decoupled weight decay with per-leaf exclusion, L2-SP anchors, and gradient accumulation.

🤖

Kaimon in the loop

The KaimonGate extension registers tools with your session: an agent drives training, validation, evaluation, prediction, and export, with runs executing in the background so no tool call ever hits a deadline.

What it is

ReactantNitro.jl is a training framework for Julia, built on Reactant.jl and Lux.jl. You declare an experiment as a struct plus a handful of hooks; the framework owns the compiled programs, the device transfers, the optimizer, the schedules, and the run's lifecycle.

The central object is a Nitro. Nitro(e) runs the setup sequence and nothing else, so evaluation and serving never depend on a train! having happened in the process; a run is what happens when you train! one.

Every field on the experiment carries one of three markers, and the marker is the point. A GraphConst field bakes into the compiled graph as a constant; a Device field becomes a traced input; an unmarked field is Host, driver-only and invisible to the tracer.

Why ReactantNitro?

Lux.jl already has a training loop, so why a framework on top of it? The plain answer is that this framework is Reactant-first, and Lux currently does not support certain things a Reactant-first training stack needs, such as gradient accumulation and the sort of phase system that frameworks like PyTorch Lightning have. ReactantNitro builds on top of Lux.jl the way Lightning builds on PyTorch, with Reactant/XLA as the accelerator under the hood.

Philosophy

This is a framework for the typical deep learning workflow: train, validate, test, predict, with dataloaders, and with Reactant compiling to an accelerator so the loop is fast. Within that context the framework tries not to make decisions that reduce what a user can express; it is not trying to be general beyond it. What it always provides is the workflow itself: the training loop, checkpointing, and a Lux.jl based model, usually a neural network.

The direct inspiration is frameworks like PyTorch Lightning. The focus, above everything else, is being the best machine learning framework for Reactant specifically.

The REPL and the Revise workflow

The package is meant to be driven from a REPL with Revise loaded. The rules are short.

• A Nitro is a fixed point. Revise a hook in the REPL and build a new handle to pick it up; the rebuild is cheap because the compile cache is module-level. The loop is on the Experiments page.
• A Device value is the one thing that changes on a live handle, via set_device!, provably without recompiling.
• The fixed-config report names the hooks you have redefined since the handle was built, so a stale handle is never silent.
• ReactantNitro.cache_stats() is the acceptance check that a change did not recompile; it is one of the tools the Recompilation page builds on.

Start here

• The Tutorial: an MNIST run from configuration to prediction, end to end.
• Pitfalls: the Reactant-first snags the framework takes care of, and why none of them raise.
• Experiments: the hook contract, the three markers, and the Revise workflow.
• Recompilation: the compile cache and when a change costs a compile.
• Optimization: parameter groups, decay, and clipping.
• Schedules: what varies with the step.
• Metrics: (sum, count) pairs, host or device residency, and finalize_metrics.
• Manual training: owning the optimizer step, GANs and beyond.
• Logging: the ten verbs, the JSON default, and the TensorBoard extension.
• Export: the wire contract and the ReactantServer bundle.
• Kaimon: driving runs from a Kaimon-hosted session, the nitro_* tools.
• The API: the docstrings, collected automatically.