{"slug": "replacing-pytorch-with-julia-and-reactant-one-year-in", "title": "Replacing PyTorch with Julia and Reactant: One Year In", "summary": "A developer reports that after nearly a year of using the Julia programming language, including the Reactant compiler, they replaced PyTorch for a machine learning workflow, citing Revise.jl's edit-without-restart REPL, automatic differentiation of plain Julia code, garbage collection with tasks and channels, and multiple dispatch. The write-up presents a PathMNIST experiment configuration using a convnextv2_tiny_fcmae_ft_in22k_in1k backbone with 9 tissue classes, a 32-dimension augmented ODE state, 100 RK4 steps, batch size 16 with 2 gradient accumulation steps for an effective batch of 32, and 5 max epochs at a 1e-3 peak learning rate on RAdam, sized to fit a 12 GB consumer GPU. The author acknowledges dynamic typing as a real trade-off but argues Julia's correctness issues are overstated for scientific computing deployments.", "body_md": "\n\n```\n@experiment struct PathMNIST\n    \"Luximm variant: selects both the architecture and the pretrained checkpoint.\"\n    variant::GraphConst{Symbol} = :convnextv2_tiny_fcmae_ft_in22k_in1k\n    \"PathMNIST tissue classes.\"\n    num_classes::GraphConst{Int} = 9\n    \"Per-channel mean the pretrained weights expect (ImageNet).\"\n    pixel_mean::GraphConst{NTuple{3, Float32}} = (0.485f0, 0.456f0, 0.406f0)\n    \"Per-channel standard deviation the pretrained weights expect (ImageNet).\"\n    pixel_std::GraphConst{NTuple{3, Float32}} = (0.229f0, 0.224f0, 0.225f0)\n    \"Augmented ODE state dimensions, predicted from the pooled features; the state also\n    carries one logit dimension per class, starting at zero.\"\n    aug_dims::GraphConst{Int} = 32\n    \"Hidden width of the ODE vector field (encoder, tanh, decoder).\"\n    node_hidden::GraphConst{Int} = 256\n    \"RK4 steps over t in [0, 1].\"\n    node_steps::GraphConst{Int} = 100\n    \"Weight of the kinetic-energy penalty on the ODE's trajectory (Finlay et al., 2020).\n    Off by default; a traced input, so turning it on or scheduling it recompiles nothing.\"\n    kinetic_weight::Device{Float32} = 0.0f0\n\n    \"Where the archive comes from on first use (MedMNIST+ on Zenodo).\"\n    data_url::String = \"https://zenodo.org/records/10519652/files/pathmnist_224.npz?download=1\"\n    \"Directory for the converted, memory-mapped .u8 files, relative to the working directory.\"\n    data_dir::String = joinpath(\"data\", \"pathmnist_224\")\n    \"Seeds the shuffle; Nitro also reads it as the run seed.\"\n    seed::Int = 0\n    \"Per-step batch: fits a 12 GB consumer GPU.\"\n    batch_size::Int = 16\n    \"Gradient accumulation steps (read by Nitro by name): an effective batch of 32.\"\n    accum::Int = 2\n    max_epochs::Int = 5\n\n    \"Peak learning rate of the one-cycle schedule (RAdam).\"\n    max_lr::Float64 = 1.0e-3\n    \"Learning rate of the NeuralODE vector field relative to the rest of the model.\"\n    node_lr_ratio::Float64 = 1.0\n    \"Fraction of the run spent warming up to max_lr.\"\n    pct_start::Float64 = 0.3\n    \"Decoupled weight decay; the NeuralODE vector field is not decayed.\"\n    weight_decay::Float64 = 1.0e-4\n\n    \"Augmentation: random 90° rotations (k = 0:3).\"\n    aug_rot90::Bool = true\n    \"Augmentation: random mirror. With aug_rot90, covers all 8 orientations.\"\n    aug_flip::Bool = true\nend\n```\n\n# Replacing PyTorch with Julia and Reactant: One Year In\n\nNearly a year ago, I decided to give the [Julia](https://julialang.org) 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](https://yuri.is/not-julia/)” 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.\n\nThese were the parts of the language that sold me:\n\n- A REPL workflow with [Revise.jl](https://github.com/timholy/Revise.jl) , where edits to code take effect in the running session without a restart\n- Julia code can be differentiated and optimized as written, without porting it to a framework’s DSL. Code from one of my ODE models, built on a cubic spline from [DataInterpolations.jl](https://github.com/SciML/DataInterpolations.jl) , drops straight into a[NonlinearSolve](https://github.com/SciML/NonlinearSolve.jl) problem with no changes\n- A modern runtime in the spirit of Go: garbage collection, plus [tasks and channels](https://docs.julialang.org/en/v1/manual/asynchronous-programming/) for concurrency\n- Multiple dispatch, which lets separately written packages work together and lets you add specializations to other people’s functions, making them faster or better integrated with your code, without opening a pull request. With agents producing more code to review than ever, I think this is a big advantage for Julia\n- Fast code that stays generic and readable\n- [SciML](https://sciml.ai) , a mature ecosystem for differential equations, nonlinear solvers and sensitivity analysis\n\nThis post is about my journey, mostly focusing on the ML aspect, but also delving into my journey through the “production” side of things.\n\n## Frustration with PyTorch\n\nWe 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:\n\n1. [torchscript](https://docs.pytorch.org/docs/stable/jit.html) 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.\n2. General-purpose model servers were slow to move to [`torch.export`](https://docs.pytorch.org/docs/stable/export.html) :[NVIDIA Triton Server](https://github.com/triton-inference-server/server) only got[non-beta support](https://github.com/triton-inference-server/pytorch_backend/pull/207) for torch.export ~3 years after torch.export was introduced, and[TorchServe](https://github.com/pytorch/serve) , PyTorch’s own server, was[archived](https://github.com/pytorch/serve/pull/3395) without ever supporting it natively.\n3. [`torch.compile`](https://docs.pytorch.org/docs/stable/generated/torch.compile.html) was underwhelming for a long time: most of the performance gains seemed[locked to a specific GPU model](https://pytorch.org/get-started/pytorch-2.0/) , and for several years I measured close to 0% improvement on everything I threw at it, even with[`fullgraph=True`](https://docs.pytorch.org/docs/stable/generated/torch.compile.html) 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.\n\nTogether, 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.\n\n## One year ago: the Julia ML ecosystem\n\nWhile I initially found a robust array of packages for scientific computing, I could immediately tell fewer people were doing “[PyTorch](https://pytorch.org)” 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:\n\n1. [Lux.jl](https://github.com/LuxDL/Lux.jl) /[Flux.jl](https://github.com/FluxML/Flux.jl) /[Zygote.jl](https://github.com/FluxML/Zygote.jl) /[CUDA.jl](https://github.com/JuliaGPU/CUDA.jl) exist and seem to be the most comparable to what I get working with PyTorch.\n2. There didn’t seem to be an equivalent to [timm](https://github.com/huggingface/pytorch-image-models) where you can just load up whatever pretrained backbone you want either.\n3. People were talking about something called Reactant, but I didn’t really understand what it was or why I needed it\n\nThe 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](https://github.com/EnzymeAD/Reactant.jl) and [Enzyme](https://enzyme.mit.edu) being the next step forward for ML in Julia.\n\n## Reactant and Enzyme: compiling and differentiating Julia\n\nReactant compiles Julia code into an [MLIR](https://mlir.llvm.org) program, Enzyme differentiates that program at the compiler level ([Moses and Churavy 2020](#ref-moses2020enzyme)), and [XLA](https://openxla.org/xla) turns the result into fused CPU/GPU/TPU kernels. [StableHLO](https://openxla.org/stablehlo) is a dialect of MLIR used by XLA, and it has a number of benefits such as promising [five years of backwards compatibility](https://github.com/openxla/stablehlo/blob/main/docs/compatibility.md). 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.\n\nOn the Julia Dispatch podcast, William Moses, who leads Enzyme and Reactant, and Avik Pal, the creator of Lux, tell [how Reactant came to be](https://podcasts.apple.com/us/podcast/julia-dispatch/id1832969355?i=1000772362563) in their own words.\n\nWith 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.\n\nTo 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`](https://pypi.org/project/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.\n\n## Betting on Reactant\n\nThe 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](https://github.com/rtqichen/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.\n\nThen 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.\n\n## What I built\n\n### ReactantNitro.jl\n\n[ReactantNitro.jl](https://github.com/EnzymeAD/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](https://github.com/timholy/Revise.jl) first workflow: you can hot reload your data processing code, sweep hyperparameters, etc. without restarting your REPL or re-compiling.\n\nWhy 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](https://enzymead.github.io/ReactantNitro.jl/stable/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.\n\nThe second is a training workflow built around Revise. In October 2025 I [wrote on Hacker News](https://news.ycombinator.com/item?id=45520854) 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](#working-with-agents), call the same entry points.\n\nNitro 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](#ref-foret2021sam); [Kwon et al. 2021](#ref-kwon2021asam)). 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.\n\n### ReactantServer.jl\n\nReactantServer 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.\n\n[ReactantServer.jl](https://github.com/EnzymeAD/ReactantServer.jl) uses Reactant’s [PJRT](https://openxla.org/xla/pjrt) interface to serve [StableHLO](https://openxla.org/stablehlo) produced by `Reactant.@compile` using the [KServe V2 protocol](https://kserve.github.io/website/latest/modelserving/data_plane/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).\n\n### Stepping stones\n\nNeither of these got built overnight. Three smaller projects came first, and ReactantNitro and ReactantServer both depend on them.\n\n[Luximm.jl](https://github.com/csvance/Luximm.jl) does for Lux.jl what [timm](https://github.com/huggingface/pytorch-image-models) 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.\n\n[gRPCClient.jl](https://github.com/JuliaIO/gRPCClient.jl) and [gRPCServer.jl](https://github.com/JuliaIO/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.\n\n## Walkthrough: from training to serving\n\nWhat 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](https://github.com/csvance/PathMNISTReactant.jl), 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.\n\n## The problem: PathMNIST\n\nWe 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.\n\n## Setting up the environment\n\n```\njuliaup add 1.13.1\ngit clone https://github.com/csvance/PathMNISTReactant.jl\ncd PathMNISTReactant.jl\njulia +1.13.1 --project=.\nusing Pkg\nPkg.instantiate()\n```\n\n## ReactantNitro.jl: Declaring the `@experiment`\n\nReactantNitro 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.\n\nEach 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.\n\n`@experiment`\n## The model\n\nWe 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](#ref-chen2018neuralode); [Dupont et al. 2019](#ref-dupont2019anode)).\n\nBecause 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:\n\nPress 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.\n\n`rk4` and `node_head`: RK4 integration and the NeuralODE head\n\n```\n# Fixed-step RK4 over t in [0, 1]. Also returns each sample's kinetic energy,\n# ∫|f(u)|² dt, from the four stage evaluations.\nfunction rk4(f, u, steps::Integer)\n    T = eltype(u)\n    dt = one(T) / T(steps)\n    half, sixth = dt / T(2), dt / T(6)\n    sq(k) = sum(abs2, k; dims = 1)\n    energy = zero(sq(u))\n    for _ in 1:steps\n        k1 = f(u)\n        k2 = f(u .+ half .* k1)\n        k3 = f(u .+ half .* k2)\n        k4 = f(u .+ dt .* k3)\n        u = u .+ sixth .* (k1 .+ T(2) .* k2 .+ T(2) .* k3 .+ k4)\n        energy = energy .+ sixth .* (sq(k1) .+ T(2) .* sq(k2) .+ T(2) .* sq(k3) .+ sq(k4))\n    end\n    return (; u, energy)\nend\n\n# Augmented NeuralODE head. The state is [a; z]: `a` comes from the pooled\n# features, and `z` starts at zero and ends as the logits.\nfunction node_head(e::PathMNIST, width::Integer)\n    # captured, not passed to @compact (they would land in st)\n    steps, A, K = e.node_steps, e.aug_dims, e.num_classes\n    D = A + K\n    return @compact(\n        proj = Dense(width => A),\n        enc = Dense(D => e.node_hidden, tanh),\n        dec = Dense(e.node_hidden => D; init_weight = zeros32, init_bias = zeros32),\n    ) do pooled                            # (width, N)\n        a0 = proj(pooled)\n        u0 = vcat(a0, zeros(eltype(a0), K) .* a0[1:1, :])   # z(0) = 0\n        (; u, energy) = rk4(u -> dec(enc(u)), u0, steps)\n        motion = sqrt.(sum(abs2, u .- u0; dims = 1) ./ sum(abs2, u0; dims = 1))\n        @return (; logits = u[(A + 1):D, :], energy, motion)\n    end\nend\n```\n\n`build_model`: the backbone, the head, and pretrained weights\n\n```\nfunction ReactantNitro.build_model(e::PathMNIST, rng)\n    # num_classes = 0 drops the ImageNet classifier and returns the last\n    # feature map. prefix = (:backbone,) points the weight loader at the\n    # @compact slot of the same name.\n    backbone, load_backbone = create_pretrained(e.variant; num_classes = 0, prefix = (:backbone,))\n    width = last(feature_info(e.variant).channels)          # 768 for Tiny\n    # Captured by the closure rather than passed to @compact, which would\n    # make them trainable parameters.\n    μ = reshape(collect(e.pixel_mean), 1, 1, 3, 1)\n    σ = reshape(collect(e.pixel_std), 1, 1, 3, 1)\n\n    model = @compact(\n        backbone = backbone,\n        norm = LayerNorm((width,)),\n        head = node_head(e, width),\n    ) do x                                 # (W, H, 3, N) UInt8, as stored on disk\n        x = to_float32(x)                  # UInt8 to [0, 1] on device: 4x less to transfer\n        x = (x .- μ) ./ σ\n        feats = backbone(x)                # (W/32, H/32, width, N)\n        pooled = dropdims(mean(feats; dims = (1, 2)); dims = (1, 2))\n        @return head(norm(pooled))         # (num_classes, N) logits\n    end\n\n    ps, st = Lux.setup(rng, model)\n    # Load here, not later: Nitro snapshots the weights build_model returns\n    # as the initial point (w0).\n    ps, st = load_backbone(ps, st)\n    return model, ps, st\nend\n```\n\n## The data pipeline\n\nOnly 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.\n\n| Augmentation | Range / distribution | Justification | \n|---|---|---|\n| 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 | \n| Mirror ( `aug_flip` ) | $\\mathrm{Bernoulli}(0.5)$ | With rotation, covers all 8 symmetries of the square uniformly; also an index permutation | \n\n`DataLoader` vs ReactantNitro\nReactantNitro supports parallel dataloaders, but there are some key differences from how that works in PyTorch.\n\n|  | PyTorch `DataLoader` | ReactantNitro + MLUtils `DataLoader` | \n|---|---|---|\n| Unit of parallelism | Worker processes ( `num_workers` ) | Tasks on the default thread pool, inside the training process | \n| Dataset in workers | A copy per worker process | One object, one memory map, shared | \n| Batch hand-off | Tensors in shared memory, handles through queues | Same address space; batches passed directly | \n| Throughput knob | `num_workers` | `julia -t N,1` (N tasks) | \n| Read-ahead | `prefetch_factor` batches per worker (default 2) | `host_batches = 2 × N` on the host,`device_batches = 1` on the device | \n| Host to device | Your loop: `pin_memory` ,`.to(device, non_blocking=True)` | One transfer task, only the fields the hooks declare | \n| Worker lifetime | Restarted each epoch unless `persistent_workers=True` | Tasks, cheap to start every epoch | \n| Short final batch | `drop_last` , one setting for the loader | Train must drop it (fixed program shape); eval keeps it, padded by Nitro | \n| Memory management | Per-process; workers never pause the training process | Shared GC; data-path allocation can pause training | \n\n## Training with ReactantNitro\n\nThe 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](#ref-finlay2020rnode)) when `kinetic_weight` is turned up. `train_metrics` runs inside the gradient program, so the dynamics are logged every step at no extra cost.\n\n`forward`, `loss` and `train_metrics`: the training hooks\n\n```\nReactantNitro.forward(e::PathMNIST, model, ps, st; img) = Lux.apply(model, img, ps, st)\n\nfunction cross_entropy(e::PathMNIST, logits, label)\n    onehot = reshape(1:e.num_classes, :, 1) .== reshape(label, 1, :)\n    return -mean(sum(onehot .* logsoftmax(logits; dims = 1); dims = 1))\nend\n\nReactantNitro.loss(e::PathMNIST, out; label) =\n    cross_entropy(e, out.logits, label) + e.kinetic_weight * mean(out.energy)\n\nReactantNitro.train_metrics(e::PathMNIST, out) =\n    (; kinetic = mean(out.energy), motion = mean(out.motion))\n```\n\n`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.\n\n`metrics` and `finalize_metrics`: macro accuracy, ECE, and the ODE’s motion\n\n``` js\nconst ECE_BINS = 15    # confidence bins for the expected calibration error\n\nfunction ReactantNitro.metrics(e::PathMNIST, out; label)\n    n = length(label)\n    probs = softmax(out.logits; dims = 1)\n    confusion = zeros(Int, e.num_classes, e.num_classes)       # rows true, columns predicted\n    samples, conf, correct = zeros(ECE_BINS), zeros(ECE_BINS), zeros(ECE_BINS)\n    for (p, l) in zip(eachcol(probs), label)\n        c, pred = findmax(p)\n        confusion[l, pred] += 1\n        b = clamp(ceil(Int, c * ECE_BINS), 1, ECE_BINS)\n        samples[b] += 1\n        conf[b] += c\n        correct[b] += pred == l\n    end\n    return (;\n        confusion = confusion => :sum,\n        calib = (; samples, conf, correct) => :sum,\n        kinetic = sum(out.energy) => n,\n        motion = sum(out.motion) => n,\n    )\nend\n\nfunction ReactantNitro.finalize_metrics(e::PathMNIST, acc, split)\n    c = acc.confusion\n    recall = [c[k, k] / sum(c[k, :]) for k in axes(c, 1) if sum(c[k, :]) > 0]\n    (; samples, conf, correct) = acc.calib\n    return (;\n        macro_acc = sum(recall) / length(recall),            # mean per-class recall\n        acc = sum(c[k, k] for k in axes(c, 1)) / sum(c),\n        ece = sum(abs.(conf .- correct)) / sum(samples),\n        acc.motion,\n        acc.kinetic,\n    )\nend\n```\n\nThe 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](#ref-smith2019superconvergence)) from [ParameterSchedulers.jl](https://github.com/FluxML/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.\n\n```\nusing ParameterSchedulers: OneCycle\n\nReactantNitro.param_group(::PathMNIST, ks) = (:enc in ks || :dec in ks) ? :node : :default\n\nReactantNitro.learning_rate(e::PathMNIST) = Float32(e.max_lr)\nReactantNitro.learning_rate(e::PathMNIST, ::Val{:node}) = Float32(e.max_lr * e.node_lr_ratio)\nReactantNitro.lambda(e::PathMNIST) = Float32(e.weight_decay)\nReactantNitro.lambda(::PathMNIST, ::Val{:node}) = 0.0f0\n\n# No decay on biases or norm scales and shifts: any leaf with at most one\n# dimension larger than 1, whatever singleton dimensions Lux gives it\nReactantNitro.no_decay(::PathMNIST, ks, x) = count(>(1), size(x)) <= 1\n\nReactantNitro.schedules(e::PathMNIST) = (;\n    eta = total -> OneCycle(\n        total, Float32(e.max_lr); percent_start = e.pct_start,\n        startval = Float32(e.max_lr / 25), endval = Float32(e.max_lr / 25 / 1.0e4)\n    ),\n)\n\n# Clip the global gradient norm loosely: it only catches the spikes that\n# occasionally make the warm-up diverge\nReactantNitro.gradient_clip_norm(::PathMNIST) = 10.0f0\n\n# Validation, checkpoint selection and export read this average of the weights\nReactantNitro.ema(::PathMNIST) = EMA(0.999)\n\n# The metric the run is judged by: checkpoint selection and `restore_best` follow it\nReactantNitro.monitor(::PathMNIST) = :macro_acc => :max\n\n# Report the loss on validation and test from its own small device program\nReactantNitro.eval_loss(::PathMNIST) = true\n```\n\nThe 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!`:\n\n```\nusing PathMNISTReactant, ReactantNitro\n\nnitro = Nitro(PathMNIST(); run_dir = \"runs/pathmnist\", restore_best = true)\ntrain!(nitro)    # ends by loading the best checkpoint by validation macro accuracy\n```\n\nFive 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:\n\n`nitro` after training: the run summary\n\n```\nnitro\nNitro for PathMNIST\n┌────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────┐\n│ phase          Done (completed)                                                                                                        │\n│ epoch          5 / 5                                                                                                                   │\n│ step           14060 / 14060                                                                                                           │\n│ params         27,912,393 in 2 groups                                                                                                  │\n│ batch_size     16                                                                                                                      │\n│ devices        1  (no mesh)                                                                                                            │\n│ elapsed        31m 53s  (this `train!` call)                                                                                           │\n│ weights        trained here, then restored to the best checkpoint, epoch 5 (`restore_best`)                                            │\n│ run_dir        runs/pathmnist   run 03JVJQ8j                                                                                           │\n│ seed           0   accum 2                                                                                                             │\n├────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────┤\n│ checkpoint                                                                                                                             │\n├────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────┤\n│ selected       epoch 5, macro_acc 0.9964974637049901, run 03JVJQ8j  (the handle holds these weights)                                   │\n│ path           runs/pathmnist/epoch-0005-step-14060-macro_acc=0.996497.jld2                                                            │\n├────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────┤\n│ metrics  (validation, epoch 5)                                                                                                         │\n├────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────┤\n│ macro_acc      0.9964974637049901                                                                                                      │\n│ acc            0.9965013994402239                                                                                                      │\n│ ece            0.0011982971885022237                                                                                                   │\n│ motion         0.27984303f0                                                                                                            │\n│ kinetic        462.27524f0                                                                                                             │\n│ loss           0.012517404f0                                                                                                           │\n├────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────┤\n│ data                                                                                                                                   │\n├────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────┤\n│ split          resolved                                                                                                                │\n│ train          5624 batches; prefetch: 8 workers, 1 on device, 16 on host                                                              │\n│ val            626 batches; prefetch: 8 workers, 1 on device, 16 on host                                                               │\n│ test           449 batches; prefetch: 8 workers, 1 on device, 16 on host                                                               │\n├────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────┤\n│ bindings                                                                                                                               │\n├────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────┤\n│ binding        what                                                                                    source                          │\n│ gradient clip  global norm 10.0; optimizer program only, and changing it recompiles it                 [method on the experiment type] │\n│ monitor        macro_acc (max); followed by checkpointer                                               [method on the experiment type] │\n│ eta            optimizer field, all groups, per-group ratio applied                                    [schedules(e)]                  │\n├────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────┤\n│ parameter groups (G = 2)                                                                                                               │\n├────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────┤\n│ group          settings                                                                                params                          │\n│ :default       eta 0.001 (x1.0), lambda 0.0001 toward zero, RAdam + Decay                              27,891,104                      │\n│ :node          eta 0.001 (x1.0), lambda 0.0 toward zero, RAdam                                         21,289                          │\n└────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────┘\n  ask it for more with `history`, `experiment`, `parameters`, `states`, `logger_info`\n```\n\n`history(nitro)`: the per-epoch series\nValidation 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:\n\n```\nhistory(nitro)\nhistory of PathMNIST  (runs/pathmnist)\n┌────────────────────────────────────────────────────────────────────────────────────┐\n│ epoch   step  train_loss  macro_acc     acc       ece  motion  kinetic     loss    │\n├────────────────────────────────────────────────────────────────────────────────────┤\n│     1   2812     0.51133     0.9844  0.9848  0.008577  0.2795    83.12  0.06029    │\n│     2   5624     0.17083     0.9867  0.9867  0.002909  0.2732   230.71  0.04378    │\n│     3   8436     0.10191     0.9913  0.9912  0.002105  0.2747   399.12  0.03040    │\n│     4  11248     0.04444     0.9953  0.9953  0.001116  0.2790   461.16  0.01802    │\n│     5  14060     0.01229     0.9965  0.9965  0.001198  0.2798   462.28  0.01252  * │\n└────────────────────────────────────────────────────────────────────────────────────┘\n  * best macro_acc (max), the checkpointer's metric\n```\n\n## Evaluating on the test set\n\nValidation 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:\n\n```\nevaluate(nitro; split = :test)\n(macro_acc = 0.9503893775255273, acc = 0.9608635097493036, ece = 0.021702159289530062, motion = 0.29635245f0, kinetic = 457.93463f0, loss = 0.14936829f0)\n```\n\n| Test metric | This run (seed 0) | Five seeds (mean ± std) | \n|---|---|---|\n| Macro accuracy | 95.04% | 95.55 ± 0.47% | \n| Accuracy | 96.09% | 96.48 ± 0.44% | \n| Expected calibration error (15 bins) | 2.17% | 1.92 ± 0.27% | \n| Loss | 0.149 | 0.136 ± 0.020 | \n\nThe 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.\n\nFor 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](#ref-kather2019)). 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](#ref-doerrich2025medmnistplus)). 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](#ref-lei2025shazam)). This model’s 95.55% macro accuracy and 96.48% accuracy, averaged over five seeds, are within that range.\n\nI 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](#ref-ignatov2024nctcrche)). 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.\n\nPer 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.\n\nStroma 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.\n\nThe 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).\n\nI 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.\n\n- **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.\n- **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.\n- **Label smoothing 0.1** on top of the jitter: no better, and much worse calibrated (ECE around 8%).\n- **10 epochs** , or**a gentler backbone learning rate** : both within noise.\n\nThe 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](#ref-woo2023convnextv2)): 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.\n\nGiven 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](https://github.com/csvance/PathMNISTReactant.jl/tree/7684532), and each run’s settings and numbers are in `blog/sweep_results.jsonl`.\n\n## Exporting the model\n\nExport 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.\n\n```\nReactantNitro.export_inputs(::PathMNIST) = [ExportSpec(\"img\", UInt8, [224, 224, 3, 1])]\nReactantNitro.export_outputs(::PathMNIST) = [ExportSpec(\"logits\")]\n```\n\nThe 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.\n\n```\nusing ReactantServerExport\n\nexport_model(nitro, ReactantServerBundle(); dir = \"bundles\", name = \"pathmnist_node\",\n    batch_sizes = [1, 16])\nbundles/pathmnist_node/\n├── manifest.yaml           1.2 KB\n├── model.b1.mlir            50 KB   StableHLO, batch 1\n├── model.b16.mlir           50 KB   StableHLO, batch 16\n└── weights.safetensors     107 MB\n```\n\nThe 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.\n\n```\nexecutable_inputs:\n  - name: \"img\"\n    dtype: \"u8\"\n    shape: \"acdn\"\n    dims: { a: 224, c: 224, d: 3 }\nexecutable_outputs:\n  - name: \"logits\"\n    dtype: \"f32\"\n    shape: \"an\"\n    dims: { a: 9 }\nbatching:\n  compiled_batch_sizes: [1, 16]\n```\n\nReactant 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.\n\n## Serving with ReactantServer\n\nReactantServer 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:\n\n[podman](https://podman.io), the [NVIDIA Container Toolkit](https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/latest/index.html), and a CDI spec so podman can see the GPU. Rootless podman works.\n\n```\nsudo nvidia-ctk cdi generate --output=/etc/cdi/nvidia.yaml\npodman run --rm --device nvidia.com/gpu=0 \\\n  -v ./bundles:/var/lib/reactantserver/models:Z \\\n  -p 8001:8001 -p 8002:8002 \\\n  ghcr.io/enzymead/reactantserver@sha256:ba640d6a1b048ae9bc728a47ade971318d37fa8cfbcc6e2539e5d096b07cb0d5\n\ncurl -s localhost:8002/readyz                # 200 once the model is compiled and serving\n```\n\n## Calling the server\n\nThe 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.\n\nFor 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.\n\n``` js\nusing ReactantServerClient, Mmap\n\nconst TILE = (3, 224, 224)                 # bytes on disk are (C, W, H, N)\n\nfunction load_split(dir, split)\n    images = open(joinpath(dir, \"$(split)_images.u8\")) do io\n        Mmap.mmap(io, Array{UInt8, 4}, (TILE..., filesize(io) ÷ prod(TILE)))\n    end\n    labels = Int.(read(joinpath(dir, \"$(split)_labels.u8\"))) .+ 1   # 0-based on disk\n    return images, labels\nend\n\nstruct TileIO <: AbstractInferenceIO\n    images::Array{UInt8, 4}           # memory-mapped, (C, W, H, N)\n    logits::Matrix{Float32}           # filled in as responses arrive\nend\nBase.length(io::TileIO) = size(io.images, 4)\nReactantServerClient.item_input_bytes(::TileIO) = prod(TILE)\n\nfunction ReactantServerClient.infer_encode_chunk!(io::TileIO, r::UnitRange, slot::PoolSlot)\n    input = scratch(slot, \"img\", (224, 224, 3, length(r)), UInt8)\n    buf = pool_view(input)\n    for (k, i) in enumerate(r)        # (C, W, H) on disk to the model's (W, H, C)\n        permutedims!(view(buf, :, :, :, k), view(io.images, :, :, :, i), (2, 3, 1))\n    end\n    return input\nend\n\nfunction ReactantServerClient.infer_decode_chunk!(io::TileIO, r::UnitRange, response)\n    io.logits[:, r] = InferOutput(\"logits\", response, Float32)\n    return nothing\nend\n\nimages, labels = load_split(\"data/pathmnist_224\", \"test\")\nlogits = zeros(Float32, 9, length(labels))\n\nkserve_init()\nmodel = KServeModel(\"grpc://127.0.0.1:8001\", \"pathmnist_node\"; max_batch_size = 16)\ninfer_async(model, TileIO(images, logits))\npred = [argmax(z) for z in eachcol(logits)]\ncount(pred .== labels) / length(labels)\n```\n\nOn 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](#training-speed-against-pytorch) 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:\n\n|  | In-process ( `evaluate` ) | Served | \n|---|---|---|\n| Accuracy | 96.09% | 96.09% | \n| Macro accuracy | 95.04% | 95.04% | \n| Predictions that differ |  | 0 of 7,180 | \n| Largest probability difference |  | 0.036 | \n\nHalf 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`](https://github.com/triton-inference-server/client) or any gRPC client that implements the spec works the same way.\n\n## Training speed against PyTorch\n\nThe 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.\n\nThe 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.\n\nBoth 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.\n\n| Workload | Stack | Step | First step (compile included) | \n|---|---|---|---|\n| CNN to NeuralODE (this post’s model) | Julia, Reactant default | **88.5 ms** | 318 s | \n|  | PyTorch `max-autotune` (fastest) | 100.3 ms | 1,420 s | \n|  | PyTorch `reduce-overhead` | 107.2 ms | 147 s | \n|  | PyTorch eager | 391.6 ms | **6 s** | \n| NeuralODE alone (MLP, RK4, 100 steps) | Julia, Reactant default | **10.8 ms** | 99 s | \n|  | PyTorch `reduce-overhead` (fastest) | 16.1 ms | 6 s | \n|  | PyTorch eager | 127.1 ms | **0.4 s** | \n\nReactant, 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.\n\n## Working with agents\n\nDay 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](https://github.com/kahliburke/Kaimon.jl), by [Kahli Burke](https://slate-portfolio.pages.dev/), 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.\n\n[Revise.jl](https://github.com/timholy/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.\n\nReactantNitro 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.\n\nAny 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.\n\n`ext/MyPackageKaimonGateExt.jl`: one tool\n\n```\nmodule MyPackageKaimonGateExt\n\nusing MyPackage, KaimonGate\n\n\"\"\"\n    mypackage_status(; verbose = false) -> String\n\nReport what MyPackage is doing in this session.\n\"\"\"\nmypackage_status(; verbose::Bool = false) = MyPackage.status_report(; verbose)\n\n# Register with the session's running gate; tools merge by name with other packages' tools.\nregister() = KaimonGate.serve(tools = [KaimonGate.GateTool(\"mypackage_status\", mypackage_status)])\n\nend\n```\n\nThe 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.\n\n## What’s next\n\nThe thing I’m most looking forward to is the SciML ecosystem following Reactant, so that the existing [OrdinaryDiffEq](https://github.com/SciML/OrdinaryDiffEq.jl) and [SciMLSensitivity](https://github.com/SciML/SciMLSensitivity.jl) packages can be called inside a Reactant-compiled program, the way [Diffrax](https://github.com/patrick-kidger/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](https://github.com/SciML/OrdinaryDiffEq.jl/pull/4436) 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](https://github.com/SciML/OrdinaryDiffEq.jl/issues/4657). SciMLSensitivity can already compile the vector-Jacobian products inside its adjoints with `ReactantVJP`.\n\nI’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](https://github.com/csvance/Luximm.jl/blob/master/ext/LuximmReactantExt.jl) by computing the variance about a detached copy of the mean and subtracting an exact correction, but only inside its own [`mean_var`](https://github.com/csvance/Luximm.jl/blob/master/src/Layers/Normalization.jl): 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:\n\n```\n# Wherever a traced program takes the variance of a Float32 4-dimensional array over its\n# first three dimensions, uncorrected, about its own mean...\n@rewrite (x::TracedRArray{Float32, 4}; dims = (1, 2, 3), corrected = false) begin\n    μ = mean(x; dims)\n    var(x; dims, corrected, mean = μ)\nend => begin\n    # ...take it about a detached copy instead, and subtract the exact correction,\n    # so the backward pass skips the reduction for the zero term.\n    μ = mean(x; dims)\n    μ₀ = ignore_derivatives(μ)\n    M = length(x) ÷ length(μ)      # from the matched shape, whatever its sizes\n    var(x; dims, corrected, mean = μ₀) .- M / (M - corrected) .* abs2.(μ .- μ₀)\nend\n```\n\nBecause 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.\n\nI opened a [feature request](https://github.com/EnzymeAD/Reactant.jl/issues/3533) 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.\n\nMuch 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.\n\n## Wrapping up\n\nA 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.\n\nOn 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.\n\nGetting 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.\n\nUsing 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](https://github.com/csvance/PathMNISTReactant.jl) reproduces this walkthrough, and ReactantNitro’s [pitfalls page](https://enzymead.github.io/ReactantNitro.jl/stable/pitfalls/) covers what I learned the hard way. [ReactantNitro.jl](https://github.com/EnzymeAD/ReactantNitro.jl) and [ReactantServer.jl](https://github.com/EnzymeAD/ReactantServer.jl) are open for issues and pull requests.\n\n## Acknowledgements\n\nThanks to:\n\n- All of the [Reactant](https://github.com/EnzymeAD/Reactant.jl) ,[Enzyme](https://github.com/EnzymeAD/Enzyme) and[Enzyme-JAX](https://github.com/EnzymeAD/Enzyme-JAX) contributors, for the foundation this work sits on, and for putting up with me as I continue to learn Reactant’s internals through trial and error\n- [William Moses](https://github.com/wsmoses) , for welcoming my projects into[EnzymeAD](https://github.com/EnzymeAD)\n- [Kahli Burke](https://slate-portfolio.pages.dev/) , research scientist in the Plasma Theory and Computation group of the Nuclear Engineering and Engineering Physics Department at UW Madison, for making it easy to bring the full ML lifecycle into agentic workflows through[Kaimon.jl](https://github.com/kahliburke/Kaimon.jl)\n- [Avik Pal](https://github.com/avik-pal) , for[Lux.jl](https://github.com/LuxDL/Lux.jl)\n- The [FluxML](https://github.com/FluxML) organization, for so many essential packages, such as[Optimisers.jl](https://github.com/FluxML/Optimisers.jl) ,[ParameterSchedulers.jl](https://github.com/FluxML/ParameterSchedulers.jl) and[Functors.jl](https://github.com/FluxML/Functors.jl)\n- [Medical Metrics Inc.](https://www.medicalmetrics.com) , for funding the development of these projects and allowing them to be open sourced\n\n## References\n\n*Advances in Neural Information Processing Systems 28 (NeurIPS 2015)*, 1171–79.\n\n[https://arxiv.org/abs/1506.03099](https://arxiv.org/abs/1506.03099).\n\n*SIAM Review*59 (1): 65–98.\n\n[https://doi.org/10.1137/141000671](https://doi.org/10.1137/141000671).\n\n*Advances in Neural Information Processing Systems*31.\n\n[https://arxiv.org/abs/1806.07366](https://arxiv.org/abs/1806.07366).\n\n*Scientific Reports*15 (1).\n\n[https://doi.org/10.1038/s41598-025-92156-9](https://doi.org/10.1038/s41598-025-92156-9).\n\n*Advances in Neural Information Processing Systems*32.\n\n[https://arxiv.org/abs/1904.01681](https://arxiv.org/abs/1904.01681).\n\n*Proceedings of the 37th International Conference on Machine Learning*, Proceedings of machine learning research, vol. 119.\n\n[https://arxiv.org/abs/2002.02798](https://arxiv.org/abs/2002.02798).\n\n*International Conference on Learning Representations (ICLR)*.\n\n[https://arxiv.org/abs/2010.01412](https://arxiv.org/abs/2010.01412).\n\n*NCT-CRC-HE: Not All Histopathological Datasets Are Equally Useful*.\n\n[https://arxiv.org/abs/2409.11546](https://arxiv.org/abs/2409.11546).\n\n*PLOS Medicine*16 (1): e1002730.\n\n[https://doi.org/10.1371/journal.pmed.1002730](https://doi.org/10.1371/journal.pmed.1002730).\n\n*Proceedings of the 38th International Conference on Machine Learning (ICML)*, Proceedings of machine learning research, vol. 139: 5905–14.\n\n[https://arxiv.org/abs/2102.11600](https://arxiv.org/abs/2102.11600).\n\n*Unifying Multiple Foundation Models for Advanced Computational Pathology*.\n\n[https://arxiv.org/abs/2503.00736](https://arxiv.org/abs/2503.00736).\n\n*Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis*, SC ’21, 1–16.\n\n[https://doi.org/10.1145/3458817.3476165](https://doi.org/10.1145/3458817.3476165).\n\n*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.\n\n[https://proceedings.neurips.cc/paper/2020/file/9332c513ef44b682e9347822c2e457ac-Paper.pdf](https://proceedings.neurips.cc/paper/2020/file/9332c513ef44b682e9347822c2e457ac-Paper.pdf).\n\n*Lux: Explicit Parameterization of Deep Neural Networks in Julia*. Zenodo.\n\n[https://doi.org/10.5281/zenodo.7808903](https://doi.org/10.5281/zenodo.7808903).\n\n*Artificial Intelligence and Machine Learning for Multi-Domain Operations Applications*, Proceedings of SPIE, vol. 11006: 1100612.\n\n[https://doi.org/10.1117/12.2520589](https://doi.org/10.1117/12.2520589).\n\n*2023 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)*, 16133–42.\n\n[https://doi.org/10.1109/CVPR52729.2023.01548](https://doi.org/10.1109/CVPR52729.2023.01548).\n\n*Scientific Data*10 (1).\n\n[https://doi.org/10.1038/s41597-022-01721-8](https://doi.org/10.1038/s41597-022-01721-8).", "url": "https://wpnews.pro/news/replacing-pytorch-with-julia-and-reactant-one-year-in", "canonical_source": "https://csvance.dev/posts/2026-10-22-reactant-end-to-end/", "published_at": "2026-10-22 12:04:06+00:00", "updated_at": "2026-10-22 12:10:14.652792+00:00", "lang": "en", "topics": ["machine-learning", "ai-research", "developer-tools", "neural-networks"], "entities": ["Julia", "Reactant", "PyTorch", "Revise.jl", "DataInterpolations.jl", "NonlinearSolve", "PathMNIST", "MedMNIST+"], "also_reported_by": [], "alternates": {"html": "https://wpnews.pro/news/replacing-pytorch-with-julia-and-reactant-one-year-in", "markdown": "https://wpnews.pro/news/replacing-pytorch-with-julia-and-reactant-one-year-in.md", "text": "https://wpnews.pro/news/replacing-pytorch-with-julia-and-reactant-one-year-in.txt", "jsonld": "https://wpnews.pro/news/replacing-pytorch-with-julia-and-reactant-one-year-in.jsonld"}}