Replacing PyTorch with Julia and Reactant: One Year In

julia
reactant
machine-learning
deep-dive
Over the past year I replaced PyTorch with Julia for my ML work. This post covers the tools I built on Reactant to train, export and serve models, and walks through all of it on a real pathology dataset.
Author

Carroll Vance

Published

October 22, 2026

Nearly a year ago, I decided to give the Julia programming language a shot; not even for machine learning work initially, but as a way to speed up a production service which was CPU bound. Now I know many people are thinking “Julia has correctness issues” and “dynamic typing in mah prod”. In my experience the first is overstated; the second is a real trade-off, and one I’ve chosen to make, because the sort of things I value when deploying services and models for scientific computing are different than the sorts of things many other engineers do.

These were the parts of the language that sold me:

This post is about my journey, mostly focusing on the ML aspect, but also delving into my journey through the “production” side of things.

Frustration with PyTorch

The PyTorch logo

We greatly accelerated our main backend service by rewriting all of the performance sensitive code in Julia. After tasting the initial success I thought to myself maybe we can do something similar with our ML stack? While PyTorch is an absolutely incredible piece of software, there were a couple of pain points that made me feel like my use case wasn’t something that was going to get much support going forward:

  1. torchscript deprecation: We had come to rely on torchscript as a flexible way to deploy models. Data dependent branching and flow control were both possible, which was nice for adding postprocessing to an already trained model. It always felt much easier to just compile our models with the extra logic built in rather than try to use whatever clunky solution our serving platform offered; we loved the convenience that torchscript provided. Having triple digits worth of models tied up in a format that had received basically no support for over five years also meant we had to make a change going forward.
  2. General-purpose model servers were slow to move to torch.export: NVIDIA Triton Server only got non-beta support for torch.export ~3 years after torch.export was introduced, and TorchServe, PyTorch’s own server, was archived without ever supporting it natively.
  3. torch.compile was underwhelming for a long time: most of the performance gains seemed locked to a specific GPU model, and for several years I measured close to 0% improvement on everything I threw at it, even with fullgraph=True guaranteeing there were no graph breaks. Many of the problems with torch.compile have since been fixed, but I was living with the same eager mode performance limitations I had been living with since I first started using PyTorch.

Together, these made me want a deployment format I could count on for years and a compiler that sped up the models I was actually training. We had already adopted Julia for our backend and were building on it extensively, so the natural next step was to see whether it could handle our ML work as well.

One year ago: the Julia ML ecosystem

The Julia logo

While I initially found a robust array of packages for scientific computing, I could immediately tell fewer people were doing “PyTorch” style ML for lack of a better term. In medical imaging it’s still quite common to use CNNs as backbones are not the ceiling on most problems, believe it or not. So I surveyed the ecosystem and here is what I found:

  1. Lux.jl/Flux.jl/Zygote.jl/CUDA.jl exist and seem to be the most comparable to what I get working with PyTorch.
  2. There didn’t seem to be an equivalent to timm where you can just load up whatever pretrained backbone you want either.
  3. People were talking about something called Reactant, but I didn’t really understand what it was or why I needed it

The core proposition of Zygote seems compelling: automatic differentiation at the language level, through almost any Julia code you write. But for PyTorch-style training workloads, its speed and device memory usage weren’t competitive with PyTorch eager mode, so I decided to give Reactant a shot; there was a lot of talk around Reactant.jl and Enzyme being the next step forward for ML in Julia.

Reactant and Enzyme: compiling and differentiating Julia

The Enzyme logo

Reactant compiles Julia code into an MLIR program, Enzyme differentiates that program at the compiler level (Moses and Churavy 2020), and XLA turns the result into fused CPU/GPU/TPU kernels. StableHLO is a dialect of MLIR used by XLA, and it has a number of benefits such as promising five years of backwards compatibility. If we can compile our models to StableHLO, that means that we should have a relatively stable platform for model training and serving going forward, so I took the dive into Reactant.

NoteThe origin story of Reactant

On the Julia Dispatch podcast, William Moses, who leads Enzyme and Reactant, and Avik Pal, the creator of Lux, tell how Reactant came to be in their own words.

With Reactant, I don’t write backward passes for custom layers. Reactant traces Julia code into a small set of StableHLO operations, Enzyme’s optimization passes simplify that program, and Enzyme differentiates the optimized program, so the gradient gets the same simplifications as the forward pass. In PyTorch, a layer built from loops and indexing is often slow under autograd, and making it fast can mean writing a custom autograd.Function with its own backward. Now when I write a new layer, I only need to consider the forward pass, and Enzyme takes care of the backward pass without giving up performance.

To be fair, none of this is exclusive to Julia. JAX users can run a function through Enzyme-JAX’s optimization pipeline with the enzyme-ad package’s enzyme_jax_ir decorator, and the same package lets JAX differentiate external C++ code. Enzyme itself works on LLVM IR, so it can differentiate code from languages that compile to LLVM: C and C++ through Clang, and Rust and Fortran through their own bindings. What Reactant adds is the frontend: it traces ordinary Julia code, so an existing Julia package can become input to these compilers by being made traceable rather than rewritten against a framework’s array API.

Betting on Reactant

The next two weeks of my life were a living hell. I had no idea what I was doing with Reactant, constantly hitting tracing errors, going OoM, dealing with hour long compile times. But I persevered and kept learning from my mistakes. I had a particularly difficult problem to iterate on that I just couldn’t get PyTorch to speed up: a large backbone combined with a NeuralODE, using torchdiffeq. The model took nearly a week to converge across 4x NVIDIA A6000 (Ampere). I kept going, convinced that the optimizations provided by XLA would prove decisive for the NeuralODE part.

Then the eureka moment came: effectively a 10x speedup over what I was seeing with PyTorch eager. All that time I had spent learning the basics of Reactant had paid off, and now I was in a position to quickly iterate on the model and solved the problem later that week. For the sake of transparency, it’s worth noting that I have since revisited a similar workload with torch.compile, and got much more comparable numbers with Reactant still holding the lead, but not an “order of magnitude”. Regardless, I was sold: it was time to build a full training and serving stack on top of Reactant.

What I built

ReactantNitro.jl

The ReactantNitro.jl logo

ReactantNitro.jl is an ML framework with a training loop designed around making Reactant easier to work with and avoiding unnecessary compiling. It has a REPL + Revise.jl first workflow: you can hot reload your data processing code, sweep hyperparameters, etc. without restarting your REPL or re-compiling.

Why another framework, you say? There are two main reasons. The first is that a compiler like Reactant warrants its own framework, one that handles the pitfalls I ran into before I created ReactantNitro: datasets ending up in the trace and slowing compiles, programs recompiling when nothing they depend on has changed, stale programs after a code edit, device transfers at unexpected times, and device memory building up across batches.

The second is a training workflow built around Revise. In October 2025 I wrote on Hacker News that waiting 15+ seconds to test small changes to my PyTorch training code was rather annoying, and that “sometimes I wish we could have a training workflow similar to how Revise works.” That is how I train now. I keep one Julia session open, edit the model, loss or data pipeline in my editor, and Revise patches the change into the running session. The next run picks it up without restarting Julia or reloading packages, and Nitro recompiles only the programs the edit reaches. The @experiment declaration makes this possible by stating which settings shape the compiled program and which can change freely. Each stage of the workflow, from setting up an experiment through training, validation, evaluation, prediction and export, is a function call at the REPL, and a Ctrl-C stops a run cleanly and hands the session back. It turns out that an interface designed for a REPL is only a few steps away from one an agent can use over MCP, so Nitro’s MCP tools, covered in Working with agents, call the same entry points.

Nitro also builds in training techniques that are easy to get subtly wrong when you implement them yourself: an exponential moving average of the weights, checkpoint averaging, and sharpness-aware minimization (SAM and ASAM) (Foret et al. 2021; Kwon et al. 2021). SAM is a good example. It takes a second gradient at perturbed weights, and that has to be combined correctly with gradient accumulation, the inner optimizer and multiple devices. In Nitro the perturbation is a first-class part of the gradient step, so SAM comes at no implementation cost: turning it on is one line, with no manual optimization loop to write and none of those details to keep in your head. The same goes for an agent: trying SAM is a one-line change it can make and evaluate, instead of an optimization loop it has to write and get right, so its effort goes into the parts of the problem that matter.

ReactantServer.jl

The ReactantServer.jl logo

ReactantServer grew out of one question: can we serve more models than their weights can fit on the GPU by hot swapping them in real time? That solved a real problem for us. We have a lot of models that are called infrequently, and together they were close to no longer fitting on our pod. Beyond that, it’s a general serving platform for StableHLO.

ReactantServer.jl uses Reactant’s PJRT interface to serve StableHLO produced by Reactant.@compile using the KServe V2 protocol. It allows you to express arbitrary pre and post processing in pure Julia without having to set up a Python environment or compile any stubs, and supports hot swapping more models than you can fit in memory at PCIe speed (typically between 10-100 ms for most CNN vision models).

Stepping stones

Neither of these got built overnight. Three smaller projects came first, and ReactantNitro and ReactantServer both depend on them.

The Luximm.jl logo

Luximm.jl does for Lux.jl what timm does for PyTorch. It ports timm’s image backbones to Lux and loads their pretrained weights directly from Hugging Face Hub. It covers ResNet, SE-ResNet, BiT ResNetV2, ConvNeXt, ConvNeXt V2, VGG, ViT and CoAtNet, and parity tests check each port against the PyTorch original. The ConvNeXt V2 backbone in the walkthrough below comes from Luximm.

gRPCClient.jl and gRPCServer.jl give Julia a native gRPC stack, which ReactantServer uses to speak KServe V2. I rewrote gRPCClient.jl on top of libcurl’s HTTP/2 support. Sébastien Celles and I wrote gRPCServer.jl.

Walkthrough: from training to serving

What would a good blog post be without actually showing how you can train and serve your own models end-to-end with the Reactant stack? Yes, it’s a real thing, you can do it today if you want. All of the code is in the companion repo, so you can follow along or run it yourself. Since we are dealing with Julia, there must be a reason to add a NeuralODE somewhere. Honestly, I don’t need one: I still find these models fascinating, especially the functions and representations they learn. Let’s start with how the full stack fits together.

The Reactant ML stack, from training to serving Training runs in one Julia process: an experiment and a Lux model feed ReactantNitro, which compiles its programs through Reactant, with Enzyme differentiating the traced program, and XLA compiling for CPU, GPU, or TPU. ReactantNitro writes a model bundle through ReactantServerExport. ReactantServer serves the bundle over gRPC, and ReactantServerClient or any Triton or KServe client calls it. Train: one Julia process Serve: a container @experiment + hooks your code: one struct, a handful of hooks Lux.jl explicit-parameter model Luximm pretrained backbone ReactantNitro.jl compiles the gradient, optimizer, and eval programs data pipeline · checkpoints · export Reactant.jl traces Julia code to MLIR (StableHLO) Enzyme reverse-mode AD on the traced program XLA, through PJRT optimizes and compiles for CPU, GPU, or TPU: any XLA backend C++ libraries; no Python anywhere model bundle written by ReactantServerExport StableHLO, weights, IO spec load ReactantServer.jl KServe V2 over gRPC (:8001) compiled ahead of time, CUDA gRPC ReactantServerClient.jl or any Triton or KServe gRPC client

The problem: PathMNIST

We are pointing the stack at PathMNIST: 224x224 tissue images in 9 classes. It’s distantly related to the problems I work on, none of which have had an “ImageNet” moment with a high quality public dataset, and it’s the closest I could find to a medical imaging dataset that isn’t a toy like MNIST but is still small enough for a demo. If you know of a better one, I’d like to hear about it.

Setting up the environment

juliaup add 1.13.1
git clone https://github.com/csvance/PathMNISTReactant.jl
cd PathMNISTReactant.jl
julia +1.13.1 --project=.
using Pkg
Pkg.instantiate()

ReactantNitro.jl: Declaring the @experiment

ReactantNitro uses a small experiment macro to ensure things are wired up correctly and Reactant’s tracer doesn’t trace over things it shouldn’t, which can be a source of slow compiles. We declare an @experiment for PathMNIST. All of ReactantNitro’s hooks dispatch on this type.

Each field’s type says how it reaches the compiled program. A GraphConst field is baked into the program as a constant, so changing it recompiles; it is for values that change the graph, such as the number of classes or a layer width. A Device field is passed to the program as a traced input, so it can be swept or scheduled without recompiling; kinetic_weight is one. Unmarked fields stay on the host and never reach the compiled program, which covers everything else, such as the data paths, the batch size and the learning rate.

@experiment struct PathMNIST
    "Luximm variant: selects both the architecture and the pretrained checkpoint."
    variant::GraphConst{Symbol} = :convnextv2_tiny_fcmae_ft_in22k_in1k
    "PathMNIST tissue classes."
    num_classes::GraphConst{Int} = 9
    "Per-channel mean the pretrained weights expect (ImageNet)."
    pixel_mean::GraphConst{NTuple{3, Float32}} = (0.485f0, 0.456f0, 0.406f0)
    "Per-channel standard deviation the pretrained weights expect (ImageNet)."
    pixel_std::GraphConst{NTuple{3, Float32}} = (0.229f0, 0.224f0, 0.225f0)
    "Augmented ODE state dimensions, predicted from the pooled features; the state also
    carries one logit dimension per class, starting at zero."
    aug_dims::GraphConst{Int} = 32
    "Hidden width of the ODE vector field (encoder, tanh, decoder)."
    node_hidden::GraphConst{Int} = 256
    "RK4 steps over t in [0, 1]."
    node_steps::GraphConst{Int} = 100
    "Weight of the kinetic-energy penalty on the ODE's trajectory (Finlay et al., 2020).
    Off by default; a traced input, so turning it on or scheduling it recompiles nothing."
    kinetic_weight::Device{Float32} = 0.0f0

    "Where the archive comes from on first use (MedMNIST+ on Zenodo)."
    data_url::String = "https://zenodo.org/records/10519652/files/pathmnist_224.npz?download=1"
    "Directory for the converted, memory-mapped .u8 files, relative to the working directory."
    data_dir::String = joinpath("data", "pathmnist_224")
    "Seeds the shuffle; Nitro also reads it as the run seed."
    seed::Int = 0
    "Per-step batch: fits a 12 GB consumer GPU."
    batch_size::Int = 16
    "Gradient accumulation steps (read by Nitro by name): an effective batch of 32."
    accum::Int = 2
    max_epochs::Int = 5

    "Peak learning rate of the one-cycle schedule (RAdam)."
    max_lr::Float64 = 1.0e-3
    "Learning rate of the NeuralODE vector field relative to the rest of the model."
    node_lr_ratio::Float64 = 1.0
    "Fraction of the run spent warming up to max_lr."
    pct_start::Float64 = 0.3
    "Decoupled weight decay; the NeuralODE vector field is not decayed."
    weight_decay::Float64 = 1.0e-4

    "Augmentation: random 90° rotations (k = 0:3)."
    aug_rot90::Bool = true
    "Augmentation: random mirror. With aug_rot90, covers all 8 orientations."
    aug_flip::Bool = true
end

The model

We create a CNN backbone (provided by Luximm.jl) that maps an image to a pooled feature vector. After a LayerNorm, a linear layer projects that vector into the first 32 dimensions of the NeuralODE’s state; the other nine, one per class, always start at zero. The state then evolves over 100 steps of RK4 integration, and the nine class dimensions at the end are the logits. There is no classifier layer after the ODE, so every prediction starts out uniform and only the flow can carry it to a class, an idea that builds on augmented NeuralODEs (Chen et al. 2018; Dupont et al. 2019).

Because the logits are part of the ODE’s state, the prediction can be read at any point along the trajectory: \(p(t) = \operatorname{softmax}(z(t))\). Here is what the trained model does with 256 test tiles. A helper script in the companion repo compiles one extra program that records \(p(t)\) every two RK4 steps, then draws the flow. It places the nine classes on the corners of a regular nonagon and plots each test tile at \(\sum_k p_k(t)\,\mathrm{corner}_k\). Every trajectory starts at the center, the uniform prediction at \(t = 0\), and the flow carries it out to a class:

Press play or drag the slider to move through the integration. Tap a class to isolate its tiles, or a tile to see all nine of its probabilities at the current t.

# Fixed-step RK4 over t in [0, 1]. Also returns each sample's kinetic energy,
# ∫|f(u)|² dt, from the four stage evaluations.
function rk4(f, u, steps::Integer)
    T = eltype(u)
    dt = one(T) / T(steps)
    half, sixth = dt / T(2), dt / T(6)
    sq(k) = sum(abs2, k; dims = 1)
    energy = zero(sq(u))
    for _ in 1:steps
        k1 = f(u)
        k2 = f(u .+ half .* k1)
        k3 = f(u .+ half .* k2)
        k4 = f(u .+ dt .* k3)
        u = u .+ sixth .* (k1 .+ T(2) .* k2 .+ T(2) .* k3 .+ k4)
        energy = energy .+ sixth .* (sq(k1) .+ T(2) .* sq(k2) .+ T(2) .* sq(k3) .+ sq(k4))
    end
    return (; u, energy)
end

# Augmented NeuralODE head. The state is [a; z]: `a` comes from the pooled
# features, and `z` starts at zero and ends as the logits.
function node_head(e::PathMNIST, width::Integer)
    # captured, not passed to @compact (they would land in st)
    steps, A, K = e.node_steps, e.aug_dims, e.num_classes
    D = A + K
    return @compact(
        proj = Dense(width => A),
        enc = Dense(D => e.node_hidden, tanh),
        dec = Dense(e.node_hidden => D; init_weight = zeros32, init_bias = zeros32),
    ) do pooled                            # (width, N)
        a0 = proj(pooled)
        u0 = vcat(a0, zeros(eltype(a0), K) .* a0[1:1, :])   # z(0) = 0
        (; u, energy) = rk4(u -> dec(enc(u)), u0, steps)
        motion = sqrt.(sum(abs2, u .- u0; dims = 1) ./ sum(abs2, u0; dims = 1))
        @return (; logits = u[(A + 1):D, :], energy, motion)
    end
end
function ReactantNitro.build_model(e::PathMNIST, rng)
    # num_classes = 0 drops the ImageNet classifier and returns the last
    # feature map. prefix = (:backbone,) points the weight loader at the
    # @compact slot of the same name.
    backbone, load_backbone = create_pretrained(e.variant; num_classes = 0, prefix = (:backbone,))
    width = last(feature_info(e.variant).channels)          # 768 for Tiny
    # Captured by the closure rather than passed to @compact, which would
    # make them trainable parameters.
    μ = reshape(collect(e.pixel_mean), 1, 1, 3, 1)
    σ = reshape(collect(e.pixel_std), 1, 1, 3, 1)

    model = @compact(
        backbone = backbone,
        norm = LayerNorm((width,)),
        head = node_head(e, width),
    ) do x                                 # (W, H, 3, N) UInt8, as stored on disk
        x = to_float32(x)                  # UInt8 to [0, 1] on device: 4x less to transfer
        x = (x .- μ) ./ σ
        feats = backbone(x)                # (W/32, H/32, width, N)
        pooled = dropdims(mean(feats; dims = (1, 2)); dims = (1, 2))
        @return head(norm(pooled))         # (num_classes, N) logits
    end

    ps, st = Lux.setup(rng, model)
    # Load here, not later: Nitro snapshots the weights build_model returns
    # as the initial point (w0).
    ps, st = load_backbone(ps, st)
    return model, ps, st
end

The data pipeline

Only the training split is augmented. Validation and test tiles go through exactly as they are stored, so every evaluation, and every request to the served model later on, sees the same pixels.

Augmentation Range / distribution Justification
Rotation by multiples of 90° (aug_rot90) \(k \in \{0, 1, 2, 3\}\), uniform Tissue has no canonical orientation; an index permutation, so no interpolation and essentially free
Mirror (aug_flip) \(\mathrm{Bernoulli}(0.5)\) With rotation, covers all 8 symmetries of the square uniformly; also an index permutation

ReactantNitro supports parallel dataloaders, but there are some key differences from how that works in PyTorch.

PyTorch DataLoader ReactantNitro + MLUtils DataLoader
Unit of parallelism Worker processes (num_workers) Tasks on the default thread pool, inside the training process
Dataset in workers A copy per worker process One object, one memory map, shared
Batch hand-off Tensors in shared memory, handles through queues Same address space; batches passed directly
Throughput knob num_workers julia -t N,1 (N tasks)
Read-ahead prefetch_factor batches per worker (default 2) host_batches = 2 × N on the host, device_batches = 1 on the device
Host to device Your loop: pin_memory, .to(device, non_blocking=True) One transfer task, only the fields the hooks declare
Worker lifetime Restarted each epoch unless persistent_workers=True Tasks, cheap to start every epoch
Short final batch drop_last, one setting for the loader Train must drop it (fixed program shape); eval keeps it, padded by Nitro
Memory management Per-process; workers never pause the training process Shared GC; data-path allocation can pause training

Training with ReactantNitro

The remaining hooks are forward, loss and train_metrics. forward is the program Nitro compiles for both training and inference, and loss is the scalar Enzyme differentiates: cross-entropy, plus the kinetic-energy penalty of Finlay et al. (2020) when kinetic_weight is turned up. train_metrics runs inside the gradient program, so the dynamics are logged every step at no extra cost.

ReactantNitro.forward(e::PathMNIST, model, ps, st; img) = Lux.apply(model, img, ps, st)

function cross_entropy(e::PathMNIST, logits, label)
    onehot = reshape(1:e.num_classes, :, 1) .== reshape(label, 1, :)
    return -mean(sum(onehot .* logsoftmax(logits; dims = 1); dims = 1))
end

ReactantNitro.loss(e::PathMNIST, out; label) =
    cross_entropy(e, out.logits, label) + e.kinetic_weight * mean(out.energy)

ReactantNitro.train_metrics(e::PathMNIST, out) =
    (; kinetic = mean(out.energy), motion = mean(out.motion))

metrics runs on the host after each validation batch. Each entry is value => mode: value => n is added up over the split and divided once by the summed n, and value => :sum is only summed, which is how the confusion matrix and the calibration bins travel. A mode can cover a whole group of values, so the three calibration sums travel together under one name. finalize_metrics turns them into the numbers that matter once per epoch: macro accuracy (the mean of per-class recall, the main metric here, since the test set’s classes are not balanced), plain accuracy, the expected calibration error over 15 confidence bins, and the mean motion and kinetic energy of the ODE, so validation also shows whether the dynamics are doing anything.

const ECE_BINS = 15    # confidence bins for the expected calibration error

function ReactantNitro.metrics(e::PathMNIST, out; label)
    n = length(label)
    probs = softmax(out.logits; dims = 1)
    confusion = zeros(Int, e.num_classes, e.num_classes)       # rows true, columns predicted
    samples, conf, correct = zeros(ECE_BINS), zeros(ECE_BINS), zeros(ECE_BINS)
    for (p, l) in zip(eachcol(probs), label)
        c, pred = findmax(p)
        confusion[l, pred] += 1
        b = clamp(ceil(Int, c * ECE_BINS), 1, ECE_BINS)
        samples[b] += 1
        conf[b] += c
        correct[b] += pred == l
    end
    return (;
        confusion = confusion => :sum,
        calib = (; samples, conf, correct) => :sum,
        kinetic = sum(out.energy) => n,
        motion = sum(out.motion) => n,
    )
end

function ReactantNitro.finalize_metrics(e::PathMNIST, acc, split)
    c = acc.confusion
    recall = [c[k, k] / sum(c[k, :]) for k in axes(c, 1) if sum(c[k, :]) > 0]
    (; samples, conf, correct) = acc.calib
    return (;
        macro_acc = sum(recall) / length(recall),            # mean per-class recall
        acc = sum(c[k, k] for k in axes(c, 1)) / sum(c),
        ece = sum(abs.(conf .- correct)) / sum(samples),
        acc.motion,
        acc.kinetic,
    )
end

The optimizer is RAdam, Nitro’s default rule, with decoupled weight decay of \(10^{-4}\) on the weights, except the NeuralODE’s vector field, which gets its own parameter group with no decay, since decay would keep pulling the dynamics back toward zero. Biases and normalization scales and shifts are never decayed: a no_decay rule excludes every parameter with at most one dimension larger than 1, which covers them however many singleton dimensions Lux stores them with. The learning rate follows a one-cycle schedule (Smith and Topin 2019) from ParameterSchedulers.jl, peaking at \(10^{-3}\) at 30% of the run. A schedule key binds to an optimizer rule field, here eta, the learning rate, or to a Device field of the experiment, so kinetic_weight could be scheduled the same way without a recompile. None of these values are part of the compiled program, so changing them does not trigger a recompile. Gradients are clipped to a global norm of 10, which only catches the occasional spike: without it, one of five seeds diverged during the warm-up. Validation, checkpoint selection and the exported model use an exponential moving average of the weights, EMA(0.999), which Nitro updates after every optimizer step alongside the raw weights that training uses. monitor names the metric the run is judged by, validation macro accuracy, and checkpoint selection and restore_best follow it. eval_loss reports the loss on validation and test from its own small program on the device.

using ParameterSchedulers: OneCycle

ReactantNitro.param_group(::PathMNIST, ks) = (:enc in ks || :dec in ks) ? :node : :default

ReactantNitro.learning_rate(e::PathMNIST) = Float32(e.max_lr)
ReactantNitro.learning_rate(e::PathMNIST, ::Val{:node}) = Float32(e.max_lr * e.node_lr_ratio)
ReactantNitro.lambda(e::PathMNIST) = Float32(e.weight_decay)
ReactantNitro.lambda(::PathMNIST, ::Val{:node}) = 0.0f0

# No decay on biases or norm scales and shifts: any leaf with at most one
# dimension larger than 1, whatever singleton dimensions Lux gives it
ReactantNitro.no_decay(::PathMNIST, ks, x) = count(>(1), size(x)) <= 1

ReactantNitro.schedules(e::PathMNIST) = (;
    eta = total -> OneCycle(
        total, Float32(e.max_lr); percent_start = e.pct_start,
        startval = Float32(e.max_lr / 25), endval = Float32(e.max_lr / 25 / 1.0e4)
    ),
)

# Clip the global gradient norm loosely: it only catches the spikes that
# occasionally make the warm-up diverge
ReactantNitro.gradient_clip_norm(::PathMNIST) = 10.0f0

# Validation, checkpoint selection and export read this average of the weights
ReactantNitro.ema(::PathMNIST) = EMA(0.999)

# The metric the run is judged by: checkpoint selection and `restore_best` follow it
ReactantNitro.monitor(::PathMNIST) = :macro_acc => :max

# Report the loss on validation and test from its own small device program
ReactantNitro.eval_loss(::PathMNIST) = true

The batch fields (img, label) are keyword arguments, so each hook declares exactly what it reads, and Nitro only transfers those fields to the device. To train, create a handle and call train!:

using PathMNISTReactant, ReactantNitro

nitro = Nitro(PathMNIST(); run_dir = "runs/pathmnist", restore_best = true)
train!(nitro)    # ends by loading the best checkpoint by validation macro accuracy

Five epochs took 32 minutes on one RTX A6000, compilation included. The first loss is 2.197, which is \(\ln 9\), because the logits and the dynamics both start at zero and the model initially predicts all nine classes equally. After training, displaying the handle prints a summary of the run: the final epoch and step, the checkpoint it restored, the last validation metrics, what each schedule and setting bound to and where it came from, and the settings each parameter group ended up with:

nitro
Nitro for PathMNIST
┌────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────┐
│ phase          Done (completed)                                                                                                        │
│ epoch          5 / 5                                                                                                                   │
│ step           14060 / 14060                                                                                                           │
│ params         27,912,393 in 2 groups                                                                                                  │
│ batch_size     16                                                                                                                      │
│ devices        1  (no mesh)                                                                                                            │
│ elapsed        31m 53s  (this `train!` call)                                                                                           │
│ weights        trained here, then restored to the best checkpoint, epoch 5 (`restore_best`)                                            │
│ run_dir        runs/pathmnist   run 03JVJQ8j                                                                                           │
│ seed           0   accum 2                                                                                                             │
├────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────┤
│ checkpoint                                                                                                                             │
├────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────┤
│ selected       epoch 5, macro_acc 0.9964974637049901, run 03JVJQ8j  (the handle holds these weights)                                   │
│ path           runs/pathmnist/epoch-0005-step-14060-macro_acc=0.996497.jld2                                                            │
├────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────┤
│ metrics  (validation, epoch 5)                                                                                                         │
├────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────┤
│ macro_acc      0.9964974637049901                                                                                                      │
│ acc            0.9965013994402239                                                                                                      │
│ ece            0.0011982971885022237                                                                                                   │
│ motion         0.27984303f0                                                                                                            │
│ kinetic        462.27524f0                                                                                                             │
│ loss           0.012517404f0                                                                                                           │
├────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────┤
│ data                                                                                                                                   │
├────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────┤
│ split          resolved                                                                                                                │
│ train          5624 batches; prefetch: 8 workers, 1 on device, 16 on host                                                              │
│ val            626 batches; prefetch: 8 workers, 1 on device, 16 on host                                                               │
│ test           449 batches; prefetch: 8 workers, 1 on device, 16 on host                                                               │
├────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────┤
│ bindings                                                                                                                               │
├────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────┤
│ binding        what                                                                                    source                          │
│ gradient clip  global norm 10.0; optimizer program only, and changing it recompiles it                 [method on the experiment type] │
│ monitor        macro_acc (max); followed by checkpointer                                               [method on the experiment type] │
│ eta            optimizer field, all groups, per-group ratio applied                                    [schedules(e)]                  │
├────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────┤
│ parameter groups (G = 2)                                                                                                               │
├────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────┤
│ group          settings                                                                                params                          │
│ :default       eta 0.001 (x1.0), lambda 0.0001 toward zero, RAdam + Decay                              27,891,104                      │
│ :node          eta 0.001 (x1.0), lambda 0.0 toward zero, RAdam                                         21,289                          │
└────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────┘
  ask it for more with `history`, `experiment`, `parameters`, `states`, `logger_info`

Validation reads the moving average of the weights, so macro accuracy is already 98.4% after the first epoch, while the learning rate is still near its peak, and it improves every epoch after that. The kinetic column shows the ODE’s dynamics growing as training goes on:

history(nitro)
history of PathMNIST  (runs/pathmnist)
┌────────────────────────────────────────────────────────────────────────────────────┐
│ epoch   step  train_loss  macro_acc     acc       ece  motion  kinetic     loss    │
├────────────────────────────────────────────────────────────────────────────────────┤
│     1   2812     0.51133     0.9844  0.9848  0.008577  0.2795    83.12  0.06029    │
│     2   5624     0.17083     0.9867  0.9867  0.002909  0.2732   230.71  0.04378    │
│     3   8436     0.10191     0.9913  0.9912  0.002105  0.2747   399.12  0.03040    │
│     4  11248     0.04444     0.9953  0.9953  0.001116  0.2790   461.16  0.01802    │
│     5  14060     0.01229     0.9965  0.9965  0.001198  0.2798   462.28  0.01252  * │
└────────────────────────────────────────────────────────────────────────────────────┘
  * best macro_acc (max), the checkpointer's metric

Evaluating on the test set

Validation comes from the same source as training (NCT-CRC-HE-100K), so it is close to saturated. The test set is CRC-VAL-HE-7K, a separate cohort of patients, so it is the better measure of how well the model generalizes. The same handle evaluates it with the programs it already compiled:

evaluate(nitro; split = :test)
(macro_acc = 0.9503893775255273, acc = 0.9608635097493036, ece = 0.021702159289530062, motion = 0.29635245f0, kinetic = 457.93463f0, loss = 0.14936829f0)
Test metric This run (seed 0) Five seeds (mean ± std)
Macro accuracy 95.04% 95.55 ± 0.47%
Accuracy 96.09% 96.48 ± 0.44%
Expected calibration error (15 bins) 2.17% 1.92 ± 0.27%
Loss 0.149 0.136 ± 0.020

The five seeds ran as one sweep in a single Julia process. The first run took 34 minutes from start to finish and the other four about 25 each, because they reused the programs the first one compiled. Every seed reaches 99.6% or 99.7% macro accuracy on validation, so the spread on the test set, from 95.0% to 96.1%, comes from how each model carries over to the new cohort. The rest of this section uses the seed 0 model, the one PathMNIST() trains by default, which is the lowest of the five on the test set.

For comparison with published results on the same test set: Kather et al., who built the dataset, reported 94.3% accuracy with a VGG19 (Kather et al. 2019). In the MedMNIST+ benchmark at 224 px, the best models reach about 96% accuracy, with 96.07% for SAM ViT-B/16 trained end to end and 96.14% for a linear probe on DINO ViT-B/16 (Doerrich et al. 2025). Large pathology foundation models land in the same place: UNI 2 and Virchow 2 both reach 95.3% balanced accuracy (the same measure as macro accuracy here) and 96.0% and 96.7% accuracy, and Shazam, which combines several of them, reaches 96.0% and 97.2% (Lei et al. 2025). This model’s 95.55% macro accuracy and 96.48% accuracy, averaged over five seeds, are within that range.

I suspect part of that ceiling comes from the dataset itself. Ignatov and Malivenko found inappropriate color normalization, JPEG artifacts that differ between classes, and corrupted tiles in NCT-CRC-HE-100K, and a classifier that only sees color histograms reached over 82% accuracy on it (Ignatov and Malivenko 2024). Some of what a model learns from the training set is therefore not tissue structure, and it may not carry over to a test cohort that was prepared differently.

A typical correct test prediction for each class: the median-confidence correct example.

A typical correct test prediction for each class: the median-confidence correct example.

Per class, seven of the nine tissues are above 94% recall. Smooth muscle and stroma are the outliers, and they are mostly mistaken for each other. The model misclassifies 281 of the 7,180 test tiles. In the confusion matrix below, each row is a true class, and each cell is the share of that class’s tiles predicted as the column’s class. The buttons switch it to raw counts, or to shares of the predicted class.

Stroma and smooth muscle are both pink, fibrous tissue made of elongated, spindle-shaped cells, and the model is often confident when it confuses them. 61 of the 281 errors have a confidence above 99%, and 10 of the 12 most confident are confusions between stroma and smooth muscle. The six most confident are shown below, one per row, each beside correctly classified tiles of the class it was predicted as and of its true class: one mucus tile predicted as adipose, and five stroma tiles predicted as smooth muscle.

The six most confident test errors, one mucus tile predicted as adipose and five stroma tiles predicted as smooth muscle, each in its own row beside correctly classified tiles of the predicted and the true class.

The six most confident test errors, one mucus tile predicted as adipose and five stroma tiles predicted as smooth muscle, each in its own row beside correctly classified tiles of the predicted and the true class.

The calibration error is small overall, and the reliability diagram shows where it comes from. 6,751 of the 7,180 test tiles land in the top bin, where confidence and accuracy agree; in every bin below it the model is overconfident, but each holds at most 119 tiles (the count is on each bar).

Reliability diagram on the test set: accuracy against confidence in 15 bins, with the number of tiles in each bin. The dashed line is perfect calibration.

Reliability diagram on the test set: accuracy against confidence in 15 bins, with the number of tiles in each bin. The dashed line is perfect calibration.

I tried a few things on an earlier version of this recipe, before the moving average and gradient clipping, all on the same 5-epoch schedule unless noted, and none of them clearly beat it. Training the same configuration with different seeds moves test macro accuracy by 0.5 to 0.8 points, so anything smaller than that is noise.

  • A wavelet texture branch: five Haar levels of subband energies, concatenated to the pooled backbone features. On RGB it changed nothing (two seeds). On haematoxylin and eosin densities, after colour deconvolution, it was ahead on two seeds and behind on the third: +0.2 points on average.
  • HSV colour jitter (hue ±0.05 of a turn, saturation ±30%, value ±20%): about 1.5 points worse on test, at 5 and at 10 epochs. It also needed gradient clipping; without it, training diverged during the warm-up.
  • Label smoothing 0.1 on top of the jitter: no better, and much worse calibrated (ECE around 8%).
  • 10 epochs, or a gentler backbone learning rate: both within noise.

The backbone may also be part of why the texture ideas added nothing. ConvNeXt V2 is pretrained as a fully convolutional masked autoencoder (FCMAE) before its ImageNet fine-tuning (Woo et al. 2023): it learns by reconstructing masked regions from their surroundings, which might make it particularly well suited to local features and textures already. SAM’s image encoder, the best end-to-end model in MedMNIST+, was also initialized from masked-autoencoder pretraining.

Given how close the published results above are to this model, I suspect we are near the limit of what this test set allows. The code for these runs is in the companion repo at commit 7684532, and each run’s settings and numbers are in blog/sweep_results.jsonl.

Exporting the model

Export turns the trained model into a bundle ReactantServer can load: StableHLO for each batch size, the weights, and a manifest describing the tensors on the wire. Two hooks define the inputs a client sends and the outputs it gets back. The model already converts UInt8 pixels to normalized Float32 inside its own program, so the wire tensor is the same img batch field forward reads, and of everything forward returns, only the logits ship; the kinetic energy and motion stay in training.

ReactantNitro.export_inputs(::PathMNIST) = [ExportSpec("img", UInt8, [224, 224, 3, 1])]
ReactantNitro.export_outputs(::PathMNIST) = [ExportSpec("logits")]

The export itself is one call on the trained handle. ReactantServerExport is a separate package that loads ReactantNitro’s export extension, so a training environment that doesn’t export never has to load it. Export is a CPU trace and does not reuse the training compile cache; for both batch sizes it took about seven minutes.

using ReactantServerExport

export_model(nitro, ReactantServerBundle(); dir = "bundles", name = "pathmnist_node",
    batch_sizes = [1, 16])
bundles/pathmnist_node/
├── manifest.yaml           1.2 KB
├── model.b1.mlir            50 KB   StableHLO, batch 1
├── model.b16.mlir           50 KB   StableHLO, batch 16
└── weights.safetensors     107 MB

The manifest describes the tensors the server expects and returns, in Julia’s axis order with the batch axis last. It also carries provenance: the full experiment configuration, the ReactantNitro and Reactant versions, and the checkpoint the weights came from.

executable_inputs:
  - name: "img"
    dtype: "u8"
    shape: "acdn"
    dims: { a: 224, c: 224, d: 3 }
executable_outputs:
  - name: "logits"
    dtype: "f32"
    shape: "an"
    dims: { a: 9 }
batching:
  compiled_batch_sizes: [1, 16]

Reactant can also call StableHLO that was produced somewhere else. Reactant.Ops.hlo_call takes a StableHLO module as text, for example one exported from PyTorch with torch_xla or from JAX, and inserts it into the traced program. It becomes part of the same program, so it is optimized together with the surrounding Julia code, and Enzyme can differentiate through it. I use this all the time to bring in PyTorch models I haven’t ported to Julia yet.

Serving with ReactantServer

ReactantServer serves every bundle in a model repository over the KServe V2 gRPC protocol. The image is pinned by digest, so the command runs exactly the server described here even after newer images are published. Mount the repository read-write, since the server stores each bundle’s compiled executables beside it so later starts skip compilation:

podman, the NVIDIA Container Toolkit, and a CDI spec so podman can see the GPU. Rootless podman works.

sudo nvidia-ctk cdi generate --output=/etc/cdi/nvidia.yaml
podman run --rm --device nvidia.com/gpu=0 \
  -v ./bundles:/var/lib/reactantserver/models:Z \
  -p 8001:8001 -p 8002:8002 \
  ghcr.io/enzymead/reactantserver@sha256:ba640d6a1b048ae9bc728a47ade971318d37fa8cfbcc6e2539e5d096b07cb0d5

curl -s localhost:8002/readyz                # 200 once the model is compiled and serving

Calling the server

The client only depends on ReactantServerClient, which doesn’t pull in Reactant or XLA, so it can run on a machine without a GPU. It reads the test split from the same memory-mapped .u8 files used for training.

For streaming a whole dataset, ReactantServerClient stages requests in a pool of buffers it allocates once and reuses, so sending the test set doesn’t allocate a new array per request or put pressure on the garbage collector. Several requests are in flight at once, and when the server is on the same host the client can send the data through shared memory instead.

using ReactantServerClient, Mmap

const TILE = (3, 224, 224)                 # bytes on disk are (C, W, H, N)

function load_split(dir, split)
    images = open(joinpath(dir, "$(split)_images.u8")) do io
        Mmap.mmap(io, Array{UInt8, 4}, (TILE..., filesize(io) ÷ prod(TILE)))
    end
    labels = Int.(read(joinpath(dir, "$(split)_labels.u8"))) .+ 1   # 0-based on disk
    return images, labels
end

struct TileIO <: AbstractInferenceIO
    images::Array{UInt8, 4}           # memory-mapped, (C, W, H, N)
    logits::Matrix{Float32}           # filled in as responses arrive
end
Base.length(io::TileIO) = size(io.images, 4)
ReactantServerClient.item_input_bytes(::TileIO) = prod(TILE)

function ReactantServerClient.infer_encode_chunk!(io::TileIO, r::UnitRange, slot::PoolSlot)
    input = scratch(slot, "img", (224, 224, 3, length(r)), UInt8)
    buf = pool_view(input)
    for (k, i) in enumerate(r)        # (C, W, H) on disk to the model's (W, H, C)
        permutedims!(view(buf, :, :, :, k), view(io.images, :, :, :, i), (2, 3, 1))
    end
    return input
end

function ReactantServerClient.infer_decode_chunk!(io::TileIO, r::UnitRange, response)
    io.logits[:, r] = InferOutput("logits", response, Float32)
    return nothing
end


images, labels = load_split("data/pathmnist_224", "test")
logits = zeros(Float32, 9, length(labels))

kserve_init()
model = KServeModel("grpc://127.0.0.1:8001", "pathmnist_node"; max_batch_size = 16)
infer_async(model, TileIO(images, logits))
pred = [argmax(z) for z in eachcol(logits)]
count(pred .== labels) / length(labels)

On the same RTX A6000, the full test set ran at about 1,000 tiles per second in batches of 16, and a single tile took 5.9 ms at the median and 6.6 ms at the 99th percentile. For comparison, a training step in the benchmark below takes 88.5 ms for 32 tiles, about 360 tiles per second. Serving is about 2.8 times faster than training, since it runs only the forward pass, without the backward pass or the optimizer update. The served model agrees with the in-process one:

In-process (evaluate) Served
Accuracy 96.09% 96.09%
Macro accuracy 95.04% 95.04%
Predictions that differ 0 of 7,180
Largest probability difference 0.036

Half the tiles agree to within \(10^{-7}\) and 99% to within \(10^{-4}\); the largest difference, on a single tile, does not change its prediction. The differences come from the server compiling and autotuning its own kernels. And because ReactantServer speaks the standard KServe V2 protocol, the client does not have to be Julia: Triton’s tritonclient or any gRPC client that implements the spec works the same way.

Training speed against PyTorch

The switch started with a 10x speedup over PyTorch eager on a CNN into a NeuralODE. A year later, I ran the comparison again on this post’s model and on the NeuralODE by itself, against both PyTorch eager and torch.compile. It only covers these two workloads: enough to show that Reactant’s performance is good and that compile time is not a deal breaker.

The biggest surprise for me was torch.compile. Coming back to it after a year, it now works smoothly on these workloads. Much of the 10x I originally saw was against PyTorch eager, and a hand-written RK4 loop compiled with mode="reduce-overhead", which captures each step in CUDA graphs, closes most of that gap. I couldn’t get torchdiffeq itself to compile, even with its RK4 solver, so both stacks here use a hand-written RK4.

Both stacks use the same weights (timm’s ConvNeXt V2 Tiny and Luximm’s port of it), batch size 16 with 2 accumulation steps, 224×224 inputs and TF32, on one RTX A6000. “Step” is the steady-state time for one training step, and the last column is the first step, including compilation.

Workload Stack Step First step (compile included)
CNN to NeuralODE (this post’s model) Julia, Reactant default 88.5 ms 318 s
PyTorch max-autotune (fastest) 100.3 ms 1,420 s
PyTorch reduce-overhead 107.2 ms 147 s
PyTorch eager 391.6 ms 6 s
NeuralODE alone (MLP, RK4, 100 steps) Julia, Reactant default 10.8 ms 99 s
PyTorch reduce-overhead (fastest) 16.1 ms 6 s
PyTorch eager 127.1 ms 0.4 s

Reactant, through XLA, is ahead of torch.compile (Inductor) on both workloads: 1.13x faster than its fastest mode on the full CNN to NeuralODE model, and 1.5x on the NeuralODE alone. Against PyTorch eager the gap is 4.4x and 12x. Compilation takes a few minutes, and ReactantNitro keeps the compiled programs for as long as the Julia process runs, so later runs reuse them unless the graph changes. PyTorch’s fastest mode on the full model, max-autotune, compiled for 24 minutes. On the NeuralODE alone, though, torch.compile is much quicker than Reactant: 6 s against 99 s.

Working with agents

The Kaimon.jl logo The Revise.jl logo

Day to day, I don’t write much code by hand anymore. Coding agents write most of it, while I review the results and make the decisions. Agents make many small edits and check each one, so the cost of starting Julia, loading packages and compiling for every check adds up quickly. Kaimon.jl, by Kahli Burke, is what makes this practical in Julia. It runs Julia sessions and exposes them to an agent over MCP, so the agent evaluates code, runs tests and searches the codebase in a warm session instead of starting a new Julia process for every command.

Revise.jl is what lets that session stay open. The agent’s edits are patched into the running session, and changes that don’t touch the compiled graph, such as the learning rate, schedules, augmentation, data loading or logging, take effect with no restart and no recompile. A run already in progress keeps running the code it started with, so edits take effect from the next run.

ReactantNitro and ReactantServer both build on this. Each has a package extension that loads when KaimonGate, the part of Kaimon that runs inside the Julia session, is present, and registers MCP tools for its own work. ReactantNitro’s tools set up an experiment; start, monitor and stop training; and validate, evaluate, predict and export (nitro_setup, nitro_train, nitro_status, nitro_stop, nitro_validate, nitro_evaluate, nitro_predict, nitro_export). ReactantServer’s tools start and stop a server, report its status and list the models it serves (rserver_start, rserver_status, rserver_models, rserver_stop). An agent can train a model, export it and serve it from the same session it uses to edit the code. Training runs in the background and the tool returns a run ID straight away, so the agent checks on the run instead of waiting for it.

Any Julia package can give agents tools the same way. It declares KaimonGate as a weak dependency and adds a package extension, and Julia compiles and loads that extension automatically whenever KaimonGate is loaded in the same session, and never otherwise, so the package itself takes on no new dependency. Inside the extension, a tool is an ordinary Julia function: KaimonGate reads its arguments and keyword arguments to build the MCP schema, and uses its docstring as the tool’s description.

module MyPackageKaimonGateExt

using MyPackage, KaimonGate

"""
    mypackage_status(; verbose = false) -> String

Report what MyPackage is doing in this session.
"""
mypackage_status(; verbose::Bool = false) = MyPackage.status_report(; verbose)

# Register with the session's running gate; tools merge by name with other packages' tools.
register() = KaimonGate.serve(tools = [KaimonGate.GateTool("mypackage_status", mypackage_status)])

end

The model in this post was built this way. An agent ran the training runs, the hyperparameter sweep behind the ablations, and the scripts that produce the figures, through Kaimon.

What’s next

The SciML logo

The thing I’m most looking forward to is the SciML ecosystem following Reactant, so that the existing OrdinaryDiffEq and SciMLSensitivity packages can be called inside a Reactant-compiled program, the way Diffrax works inside JAX. That would bring their adaptive and stiff solvers and choice of adjoint methods into compiled programs, instead of the hand-written RK4 I use today. It isn’t here yet, but work is moving toward it. A draft pull request makes OrdinaryDiffEq’s solve trace under Reactant’s @compile for explicit solvers, with the remaining pieces, such as saveat, callbacks and implicit solvers, tracked in an issue. SciMLSensitivity can already compile the vector-Jacobian products inside its adjoints with ReactantVJP.

I’m also looking forward to writing Enzyme-JAX rewrite rules in Julia itself. Luximm’s normalization layers show why. Reverse-mode AD of a variance taken about its own mean spends a full-size reduction on a term, \(\partial\sigma^2/\partial\mu\), that is always zero, because \(\mu\) is the mean. Luximm avoids it by computing the variance about a detached copy of the mean and subtracting an exact correction, but only inside its own mean_var: every other package that computes a variance this way still pays for the zero term. If the same fix could be written as a pattern and its replacement in ordinary Julia, the compiler could apply it wherever that pattern appears in a traced program, so one optimization would cover an entire ecosystem instead of a single implementation. The rule would declare the inputs it applies to the way a method signature does: the element type and rank of the traced array, and the values of dims and corrected, which become constants in the compiled program. The sizes stay open, so it matches a variance over any image or batch size, and the compiler traces both sides against the shapes it finds. Something like this, which is a sketch of what I would like rather than an API that exists:

# Wherever a traced program takes the variance of a Float32 4-dimensional array over its
# first three dimensions, uncorrected, about its own mean...
@rewrite (x::TracedRArray{Float32, 4}; dims = (1, 2, 3), corrected = false) begin
    μ = mean(x; dims)
    var(x; dims, corrected, mean = μ)
end => begin
    # ...take it about a detached copy instead, and subtract the exact correction,
    # so the backward pass skips the reduction for the zero term.
    μ = mean(x; dims)
    μ₀ = ignore_derivatives(μ)
    M = length(x) ÷ length(μ)      # from the matched shape, whatever its sizes
    var(x; dims, corrected, mean = μ₀) .- M / (M - corrected) .* abs2.(μ .- μ₀)
end

Because Enzyme differentiates the program after the rule has rewritten it, the rule only has to describe the forward computation: the cheaper backward pass comes for free, so all the effort goes into the optimization itself.

I opened a feature request for this after hearing William mention it; he confirmed it’s part of the plan for Reactant, but emphasized that work still needs to be done on the Julia integration side before it can happen.

Much of the potential with this resides in optimizations that could never become defaults because they violate certain assumptions, as a compiler’s built-in rules have to be correct for every program. The variance rule has a cheaper relative like that: detach the mean and drop the correction entirely. First derivatives stay exact, but higher derivatives become wrong, which is fine for a training loop that never differentiates twice and unacceptable for one with gradient penalties. Rules written in Julia could live in a package and be turned on only for the programs where their assumptions hold, so a set of optimizations tailored to one kind of workload could ship like any other library, without forking the compiler or changing anyone else’s results.

Wrapping up

A year ago I started using Julia to speed up a CPU-bound service. Today I train, export and serve models with it, and the bet solved each of the problems that pushed me away from PyTorch. Every PyTorch model we had converted to StableHLO, including object detection models, several BERT models and many other computer vision models. ReactantServer serves them directly, with any pre and post processing written in plain Julia. We have been testing it for months, and it is part of an upcoming regulated release.

On the workloads I measured, Reactant trains faster than torch.compile, and when I write a custom layer, I no longer need a hand-written backward pass to get acceptable performance. I haven’t trained a PyTorch model in half a year. The walkthrough in this post covers the whole path on a real dataset: a ConvNeXt V2 backbone with a NeuralODE head, trained with ReactantNitro to results in line with published models on the same test set, exported to StableHLO, and served by ReactantServer with the same predictions it makes in process.

Getting here took more than Reactant. I spent a significant amount of time getting the gRPC client and server packages stable enough for production use, ported timm’s backbones to Lux, and built my own ML framework and serving platform. The day-to-day workflow has made that effort worth it. Agents train, export and serve models through MCP, and they do well with a REPL and Revise driven workflow, where each small edit can be checked in a session that is already running.

Using Julia hasn’t impeded my work, and I’m more excited about the language’s future than ever. If you want to try it, the companion repo reproduces this walkthrough, and ReactantNitro’s pitfalls page covers what I learned the hard way. ReactantNitro.jl and ReactantServer.jl are open for issues and pull requests.

Acknowledgements

Thanks to:

References

Bengio, Samy, Oriol Vinyals, Navdeep Jaitly, and Noam Shazeer. 2015. “Scheduled Sampling for Sequence Prediction with Recurrent Neural Networks.” Advances in Neural Information Processing Systems 28 (NeurIPS 2015), 1171–79. https://arxiv.org/abs/1506.03099.
Bezanson, Jeff, Alan Edelman, Stefan Karpinski, and Viral B. Shah. 2017. “Julia: A Fresh Approach to Numerical Computing.” SIAM Review 59 (1): 65–98. https://doi.org/10.1137/141000671.
Chen, Ricky T. Q., Yulia Rubanova, Jesse Bettencourt, and David Duvenaud. 2018. “Neural Ordinary Differential Equations.” Advances in Neural Information Processing Systems 31. https://arxiv.org/abs/1806.07366.
Doerrich, Sebastian, Francesco Di Salvo, Julius Brockmann, and Christian Ledig. 2025. “Rethinking Model Prototyping Through the MedMNIST+ Dataset Collection.” Scientific Reports 15 (1). https://doi.org/10.1038/s41598-025-92156-9.
Dupont, Emilien, Arnaud Doucet, and Yee Whye Teh. 2019. “Augmented Neural ODEs.” Advances in Neural Information Processing Systems 32. https://arxiv.org/abs/1904.01681.
Finlay, Chris, Jörn-Henrik Jacobsen, Levon Nurbekyan, and Adam M. Oberman. 2020. “How to Train Your Neural ODE: The World of Jacobian and Kinetic Regularization.” Proceedings of the 37th International Conference on Machine Learning, Proceedings of machine learning research, vol. 119. https://arxiv.org/abs/2002.02798.
Foret, Pierre, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur. 2021. “Sharpness-Aware Minimization for Efficiently Improving Generalization.” International Conference on Learning Representations (ICLR). https://arxiv.org/abs/2010.01412.
Ignatov, Andrey, and Grigory Malivenko. 2024. NCT-CRC-HE: Not All Histopathological Datasets Are Equally Useful. https://arxiv.org/abs/2409.11546.
Kather, Jakob Nikolas, Johannes Krisam, Pornpimol Charoentong, et al. 2019. “Predicting Survival from Colorectal Cancer Histology Slides Using Deep Learning: A Retrospective Multicenter Study.” PLOS Medicine 16 (1): e1002730. https://doi.org/10.1371/journal.pmed.1002730.
Kwon, Jungmin, Jeongseop Kim, Hyunseo Park, and In Kwon Choi. 2021. “ASAM: Adaptive Sharpness-Aware Minimization for Scale-Invariant Learning of Deep Neural Networks.” Proceedings of the 38th International Conference on Machine Learning (ICML), Proceedings of machine learning research, vol. 139: 5905–14. https://arxiv.org/abs/2102.11600.
Lei, Wenhui, Yusheng Tan, Anqi Li, et al. 2025. Unifying Multiple Foundation Models for Advanced Computational Pathology. https://arxiv.org/abs/2503.00736.
Moses, William S., Valentin Churavy, Ludger Paehler, et al. 2021. “Reverse-Mode Automatic Differentiation and Optimization of GPU Kernels via Enzyme.” Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis, SC ’21, 1–16. https://doi.org/10.1145/3458817.3476165.
Moses, William, and Valentin Churavy. 2020. “Instead of Rewriting Foreign Code for Machine Learning, Automatically Synthesize Fast Gradients.” In Advances in Neural Information Processing Systems, edited by H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin, vol. 33. Curran Associates, Inc. https://proceedings.neurips.cc/paper/2020/file/9332c513ef44b682e9347822c2e457ac-Paper.pdf.
Pal, Avik. 2023. Lux: Explicit Parameterization of Deep Neural Networks in Julia. Zenodo. https://doi.org/10.5281/zenodo.7808903.
Smith, Leslie N., and Nicholay Topin. 2019. “Super-Convergence: Very Fast Training of Neural Networks Using Large Learning Rates.” Artificial Intelligence and Machine Learning for Multi-Domain Operations Applications, Proceedings of SPIE, vol. 11006: 1100612. https://doi.org/10.1117/12.2520589.
Woo, Sanghyun, Shoubhik Debnath, Ronghang Hu, et al. 2023. “ConvNeXt V2: Co-Designing and Scaling ConvNets with Masked Autoencoders.” 2023 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), 16133–42. https://doi.org/10.1109/CVPR52729.2023.01548.
Yang, Jiancheng, Rui Shi, Donglai Wei, et al. 2023. “MedMNIST v2: A Large-Scale Lightweight Benchmark for 2D and 3D Biomedical Image Classification.” Scientific Data 10 (1). https://doi.org/10.1038/s41597-022-01721-8.