Migrating from v2.x to v3.0

v3.0 is the largest single release since v2 — the runtime is now implemented entirely in Python (running on the PyTorch backend), with AOT-Inductor as the path to a portable compiled artifact. Most input files written for v2.1.2 continue to work unchanged; the changes below are the ones to look at first.

Why migrate?

Measured as peak throughput (material updates per ms at the optimal batch size \(N^*\) that each runtime hits its vectorisation ceiling at — see Benchmarking for the methodology), comparing v3.0 against v2.1.6 on the same node (NVIDIA RTX A5000 + Intel Xeon Gold 6346, 32 hyperthreaded cores):

  • CUDA, every scenario: v3 is 1.0× to 4.5× faster than v2 — every benchmarked model + driver combination is at least as fast on v3, most are 2-4× faster.

  • CPU, non-crystal-plasticity: v3 is 1.4× to >10× faster than v2 across elasticity, viscoplastic hardening, Chaboche (2/4/6 back stresses), GTN poroplasticity, and \(J_2\) radial-return.

  • CPU, crystal-plasticity (known regression): v3 is 3× to 30× slower than v2 on scpcoup, scpdecoup, scpdecoupexp, tcpsingle, tcprandom. v3’s CP optimisation work targeted the CUDA AOTI path (per-slip-reduction-fused Triton kernels, K-batched JVP with broadcast-state tangents); the equivalent CPU codegen is currently slower than v2’s eager TorchScript-JIT path. Migrate CP workloads to CUDA to capture the 1.3-4.1× win there; CP-on-CPU is a planned follow-up.

The full per-batch sweep, fitted vectorisation-efficiency parameters, and raw CSVs are documented in Benchmarking.

Install / consumption

pip install neml2 is now the only supported install path. The wheel bundles the runtime shared library, public C++ headers, CMake config exports, and pkg-config files alongside the Python package, so a single install serves both Python and C++ consumers — see Basic installation.

If you previously built from source against a custom LibTorch, the source-build flow runs through pip install -e . (point it at your libtorch with --config-settings=cmake.define.torch_ROOT=...); see Building from source.

Removed CMake presets

CMakePresets.json is gone — pip install -e is the single build entry point, with no preset to remember:

  • The editable install builds the runtime, the bindings, and the C++ test executables at RelWithDebInfo; a shipped wheel uses Release. Override either with --config-settings=cmake.build-type=....

  • The sanitizer / coverage flavors run through the same install (--config-settings=cmake.build-type=Coverage, …); see Building from source.

Removed CMake options

NEML2_PCH, NEML2_JSON, NEML2_CSV are gone. Precompiled headers made a measurable difference in v2 when the C++ tower had thousands of TUs; the v3 codebase has a handful, and the option no longer pays for itself. JSON and CSV support are always on.

Removed CLI tools

neml2-diagnose and neml2-time were removed. The remaining CLI surface is:

Tool

Purpose

neml2-run

Drive a model through a load history.

neml2-inspect

Print the structural summary of a model.

neml2-syntax

Browse the registered-object catalog.

neml2-compile

Export a model to an AOT-Inductor package.

neml2-stub

Regenerate .pyi stubs for the pybind11 extensions.

See CLI utilities for the full reference.

Authoring custom models

Custom models are now written as a Python class deriving from neml2.models.model.Model (which is itself a torch.nn.Module). The authoring surface is documented end-to-end in Extension — in summary:

  • Schema declaration uses the helpers in neml2.schema (input, output, parameter, …) rather than the C++ options.add_input / add_parameter family.

  • Registration uses the @register_neml2_object("TypeName") decorator rather than the register_NEML2_object C++ macro.

  • forward(self, *typed_inputs, v=None, v2=None, vh=None) replaces the C++ set_value virtual; the optional v / v2 / vh kwargs carry the first- and second-order chain-rule channels.

If you were maintaining a C++ subclass of Model against v2, the Declaring inputs, outputs, and parameters, The forward operator, and Connecting to input files tutorials walk through the equivalent Python pattern on a single running example.

Compiled-model story: TorchScript → AOTI

neml2-jit and the TorchScript path are gone. The portable artifact is now an AOT-Inductor (AOTI) package built with neml2-compile:

neml2-compile input.i --model elasticity

The output layout is per-device: neml2-compile writes the compiled kernels and a _meta.json into <output-dir>/<model>/<device>/, and places a standalone HIT stub <output-dir>/<model>_aoti.i next to the <model>/ folder. The stub points at the artifact folder via an absolute path; the loader picks the subfolder for the running device automatically. Passing --device cpu cuda emits both cpu/ and cuda/ subfolders under the same <model>/ folder, ready for the C++ multi-device dispatcher. The stub is not relocatable without recompiling (the artifact path is absolute). The artifact loads in milliseconds from either Python or C++ without re-parsing the input file. See Compiled models for the round-trip walkthrough and AOTI packages for the package-format reference.

Input-file changes

The HIT format is unchanged. Most input files from v2.1.2 work verbatim. The notable differences:

[Tensors] collapses to Python + CSV<Type>

The big input-file change is in the [Tensors] section. v2.1.2 shipped a registered catalog of tensor-constructor types — UserTensor, FullTensor, LinspaceTensor, LogspaceTensor, IdentityTensor, GaussianTensor, RandomTensor, FillR2, FillSR2, FillRot, FillWR2, Orientation, SymmetryFromOrbifold, FromTorchScript, and the per-primitive type names (Scalar, R2, SR2, …). v3 keeps only two built-in flavors:

  • Python — an inline PyTorch expression evaluated against a namespace pre-populated with torch, math, np, and every typed wrapper from neml2.types (Scalar, SR2, R2, MRP, …). Cross-references to other [Tensors] entries by name resolve lazily, so one block can build on another.

  • CSV<Type> — CSVScalar, CSVSR2, CSVVec, CSVWR2 for loading reference snapshots from disk. These are the only typed-constructor blocks that survived.

Everything else moves into a type = Python block with an expr that calls the equivalent torch primitive. The wrapper that the expression returns can chain .sub_batch.retag(...) to tag a sub-batch axis — the same source of truth sub_batch_ndim / sub_batch.expand_at used to encode separately.

# v2.1.2 — registered constructor type per shape / pattern
[Tensors]
  [times]
    type = LinspaceScalar
    start = 0
    end = 1
    nstep = 5
  []
  [strain_template]
    type = FillSR2
    values = '0.01 0 0 0 0 0'
  []
[]

# v3.0 — same data, one Python block per entry
[Tensors]
  [times]
    type = Python
    expr = 'linspace(Scalar(0).dynamic_batch, Scalar(1).dynamic_batch, 5)'
  []
  [strain_template]
    type = Python
    expr = 'SR2.fill(0.01, 0, 0, 0, 0, 0)'
  []
[]

Bare numeric / list option values (used directly inside [Models] or [Drivers] blocks) still parse the same way and don’t need a [Tensors] entry at all. The migration is only for blocks that previously named a constructor type.

Scheduler / dispatcher input wiring

The work-dispatcher is back, but C+±only and not input-file-wired. Drop the scheduler = '<name>' option and any [Schedulers] block from your input files — the v2 HIT surface is gone. In v3 dispatch is a feature of the compiled C++ runtime: neml2-compile --device cpu cuda emits one artifact per device, and neml2::aoti::load_model(stub, name, scheduler) runs a batch across them via SimpleScheduler / MPISimpleScheduler, configured in C++ source rather than the .i. See cpp-dispatch — dispatching across devices.

Python (neml2.load_model, neml2-run) stays eager and single-device; the per-device sweep that pyzag / your own loop performs remains the Python multi-device path.

Python API surface

The user-facing Python surface is mostly stable:

  • neml2.load_model(path, name) and neml2.load_input(path) work as in v2.

  • Typed wrappers (Scalar, Vec, R2, SR2, WR2, SSR4, MRP, MillerIndex) live in neml2.types as before. Their underlying tensor is .data.

  • Scalar(<number>) now accepts a plain Python number directly (defaults to torch.float64). v2 required Scalar(torch.tensor(<number>, dtype=torch.float64)).

The notable surface changes:

  • neml2.es and the assembled-vector / assembled-matrix submodule moved to neml2.es.

  • neml2.tensors (the submodule re-exporting typed wrappers from the C++ bindings) is gone; the wrappers live under neml2.types.

  • The neml2.postprocessing module is renamed neml2.texture (ODF reconstruction, pole / inverse-pole figures, IPF reductions). The crystal symmetry operators it used now come from neml2.ops.symmetry, and orientations are the neml2.types.MRP wrapper (the former Rot).

  • The neml2.reader module is gone — neml2.load_input covers the parser-facing surface.

Example notebooks

The standalone notebooks that lived under python/examples/ in v2 are re-homed as worked examples under the physics-module pages they exercise. The KWN precipitation examples are migrated — see 316H two-phase precipitation (carbide + Laves) (316H carbide + Laves) and Al–Cu TTP diagram (Al–Cu TTP diagram). The deterministic and statistical calibration notebooks are now Parameter calibration. The crystal-plasticity examples are migrated too — see Crystal plasticity for the separated-vs-integrated formulation comparison and the pole-figure / ODF reconstruction (now built on neml2.texture). The explain.ipynb LLM demo is dropped with the neml2.reader module it used.

Documentation pipeline

The documentation pipeline moved from Doxygen + custom Python scripts to Sphinx with the shibuya theme and MyST-NB. The Python contract:

pip install ".[dev]" -v
doc/scripts/build.sh

See Documentation for the contributor view of the doc build (including the --clean / --serve / --port flags exposed by the wrapper) and Building from source for the source-build view.