NEML2 exists to solve a specific and consequential problem. Nuclear reactor components, aerospace alloys, and energy-storage materials all obey constitutive laws—mathematical descriptions relating stress, strain, temperature, and time that can involve dozens of coupled physical processes operating simultaneously. Modeling these behaviors inside a full multiphysics simulation such as the MOOSE framework, which Argonne and its partners use for massively parallel reactor analysis, requires constitutive models that can be composed from small, independently written pieces. NEML2 achieves this through a registry-factory pattern: contributors register self-contained model fragments, and users assemble them into a complete material description at runtime by parsing an input file. The same pattern underpins MOOSE itself, and it is precisely what makes the library attractive to a large and distributed community—no single contributor needs to understand the entire edifice to add a new physics capability.

The trouble is that runtime flexibility and GPU efficiency pull in opposite directions. A well-optimized device kernel needs to know, before execution begins, exactly which operations will run and in what order, so the compiler can fuse them into minimal memory traffic and eliminate redundant computation. Registry-factory composition, by contrast, does not settle the structure of the computation until an input file is parsed at runtime. The previous version of NEML2, released in 2025 and itself a landmark in merging machine-learning infrastructure with materials modeling, resolved this tension with TorchScript tracing and just-in-time compilation: the runtime composition was converted into a concrete computational graph and partially compiled on the fly. But that approach fused only narrow islands of the graph, interpreting the rest. The coverage limit became a performance ceiling, and lifting it became the central motivation for the rewrite.

Version 3 answers the problem with a fundamentally different compilation philosophy. Instead of compiling fragments of the graph just in time, the new release compiles essentially the entire graph ahead of time, once, offline, using AOTInductor—a modern graph-compilation backend from the PyTorch ecosystem. A command-line tool called neml2-compile converts a Python-authored model into a portable, self-describing artifact containing per-segment compiled graphs and versioned metadata. The artifact can then be loaded from pure C++ with no Python interpreter and no compiler toolchain present on the target machine—an essential property for embedding NEML2 inside MOOSE, where introducing a Python dependency would complicate deployment on some of the world’s largest computing facilities.

The move to ahead-of-time compilation dictated a second, equally consequential change: the authoring language. PyTorch’s torch.export tracer, the machinery that records computational graphs, instruments live Python execution; a C++ forward pass is opaque to it, leaving the compiler nothing to fuse. All physics in version 3 therefore lives in Python, written as ordinary subclasses of torch.nn.Module, while the surviving C++ is generic plumbing—an artifact loader, a Newton–Krylov solver driver, a multi-device dispatcher, and a carefully designed exception taxonomy—with no constitutive-model logic of its own. The practical payoff for researchers extends well beyond performance. A scientist can now develop a custom material model as an ordinary Python package, distribute it independently, and reference it in an input file exactly as if it were built into the library. No forking the repository, no rebuilding binaries, no coordinating releases with the core team.

The architecture is more subtle than it first appears, because one category of physics resists graph compilation entirely. Data-dependent Newton iteration—the implicit solver loops at the heart of nonlinear material models—cannot be traced by torch.export, since the number of iterations is not known in advance. NEML2’s solution is to partition each model at every implicit-solve boundary into independently compiled segments, with the convergence loop running in host-side C++ around them. This splitting had a happy side effect: separating the compiled Jacobian graph from the linear solve that consumes it enabled a genuine solver layer, something version 2 lacked altogether. The new release ships direct DenseLU and Schur-complement solvers—the latter exploiting block-diagonal structure to scale linearly rather than cubically with the number of local sub-systems—alongside matrix-free Krylov methods including GMRES and BiCGStab, each with a family of preconditioners. An adaptive substepping scheme wraps any failing solve, recursively bisecting the load increment per batch row in host control flow until convergence is achieved.

The physics catalog itself has grown substantially. Eight new modules have been added since the original publication: KWN precipitation kinetics, which tracks how nanoscale precipitates nucleate and grow in alloys; an expanded crystal-plasticity catalog describing deformation at the level of individual slip systems; viscoelasticity; traction–separation laws for fracture interfaces; porous flow; phase-field fracture; chemical reaction networks; and a finite-volume engine. Each module is expressed as a set of composable primitives consistent with the library’s architecture, and end-to-end worked examples now exercise the KWN and crystal-plasticity capabilities, giving newcomers concrete templates for their own model development.

Performance testing, conducted head-to-head against the legacy just-in-time path using the same vectorization-efficiency benchmark model from the original publication, shows that ahead-of-time Inductor compilation is at least as fast as the old approach in every GPU scenario measured, with gains ranging from 1.0 to 4.5 times and no additional device memory consumption. On CPUs, the picture is more nuanced but still striking: the new pipeline wins broadly, by as much as 64 times, though it regresses by factors of 3 to 33 for crystal plasticity—a limitation the authors trace to Inductor’s CPU code generation for the per-slip-system reduction pattern characteristic of that physics. The practical guidance is clear: crystal-plasticity workloads belong on GPUs, where the new compilation still wins by 1.3 to 3.9 times. Importantly, the architectural choice is validated by independent evidence: the NequIP and Allegro machine-learning interatomic potential projects, confronting the identical problem of deploying composable, GPU-targeted models into non-Python hosts, have reported speedups of similar or greater magnitude after making the same migration from TorchScript to AOTInductor.

Robustness and maintainability received as much attention as raw speed. Every entry point in the library funnels errors through a three-type taxonomy keyed on a single recoverable predicate, allowing a host application to distinguish a Newton convergence failure—normally handled by cutting the time step and retrying—from a fatal condition requiring termination. The compiled artifact is exposed to C++ consumers behind a deliberately stable, PImpl-hidden application binary interface, so that internal refactoring of NEML2 does not force a rebuild of every host code that links against it. Developer tooling includes a unified logging system with per-channel severity levels, a separately versioned input-file parser named nmhit that replaces an external C++ dependency, a dependency manager that keeps all version pins—including the compiled artifact schema version—in lockstep, and PyPI binary wheels validated across a compatibility matrix of forty combinations of PyTorch versions, Python versions, and operating systems. In a sign of the times, the repository even ships checked-in Claude Code developer skills that encode the project’s contribution conventions, reflecting a contributor base that is now partly AI-assisted.

The MOOSE integration was overhauled in parallel. A single NEML2ModelHandle interface now abstracts over both the eager and compiled runtimes, so switching between interpreted development and compiled production deployment is a configuration change rather than a code change. Device selection—CPU versus GPU, or specific GPU indices—became an ordinary input-file parameter, and batches of material points are dispatched across multiple devices through NEML2’s own MPI-aware scheduler, the continuation of a work-dispatch architecture the team introduced in an earlier Argonne technical report. The result is a library that nuclear-energy modelers can deploy across the full spectrum of hardware, from a laptop running Python interactively to an exascale machine driving millions of coupled material points on thousands of GPUs. Sponsored by the Department of Energy’s Nuclear Energy Advanced Modeling and Simulation Program, the project is released under the MIT license, installable with a single pip command, and documented for a community whose ambition is nothing less than predictive, validated simulation of advanced reactors and the materials that must survive inside them.