CLAUDE.md

CLAUDE.md

This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository.

Overview

NEML2 is a Python-native material modeling library that vectorizes constitutive model evaluation on CPU/GPU using PyTorch as the tensor backend. Models are plain torch.nn.Module subclasses composed from small reusable pieces, the framework auto-resolves dependencies between them, and most users interact via HIT input files (the same format used by MOOSE) plus either the Python API or the neml2-run / neml2-compile CLIs.

The legacy C++ tower from v2.x was retired in the v3 migration. The C++ that remains lives under neml2/csrc/: the AOT-Inductor runtime (neml2/csrc/aoti/) that loads the .pt2 packages produced by neml2-compile, plus the work scheduler / dispatcher (neml2/csrc/dispatchers/) that spreads a batched evaluation across CPU/GPU(s). Both compile into a single shared library, neml2/lib/libneml2.so, built once as part of the wheel and invisible to most contributors. A second, separate library neml2/lib/libneml2_eager.so (neml2/csrc/eager/) embeds a CPython interpreter to run a model straight from its .i with no compile step — the fast-to-start path for downstream C++ unit tests; it is kept separate precisely so libneml2.so stays Python-free. See for the v2 → v3 rewrite summary.

The six evaluation routes

A model is authored once in Python, then evaluated through one of six runtimes. Each carries a host-mode codename — use these in commits, issues, and review. Full reference (capability matrix, when-to-use): .

Codename	Entry point	Compile	Host / library
py-eager	neml2.load_model	none	Python
py-jit	neml2.compile	in-process torch.compile	Python
py-aoti	neml2.aoti.Model	offline (neml2-compile)	Python via pybind over libneml2.so
cpp-aoti	neml2::aoti::Model	offline (neml2-compile)	C++ (libneml2.so)
cpp-dispatch	neml2::aoti::DispatchedModel	offline (neml2-compile)	C++ multi-device (libneml2.so)
cpp-eager	neml2::eager::Model	none	C++ + embedded CPython (libneml2_eager.so)

Parity is an invariant. All six expose the same forward / jvp / jacobian surface over the same authored model. A change to the model boundary, the chain rule, or the AOTI metadata contract must be carried across every affected route — a fix that lands only in py-eager silently diverges the other five. The one deliberate asymmetry: cpp-eager is plain-batch only (sub-batch models such as crystal plasticity are rejected); every other route supports sub-batch.

Hard rules — these supersede convenience

These three rules are non-negotiable across the codebase. They exist because violations are individually local but collectively cascade into the kind of silent metadata-loss bug that consumes weeks of debugging. Re-read this section before any non-trivial edit.

1. No raw torch.Tensor returns from neml2 internal functions

Every internal neml2 function, method, leaf, helper, and chain-rule operator returns a typed wrapper (Scalar, SR2, Tensor, ...) with correct sub_batch_ndim and sub_batch_labels. Internal callers should never have to re-attach metadata; the producer always returns typed.

The only legitimate raw-tensor surfaces are the framework-imposed boundaries:

• AOTI shim (neml2/aoti/, neml2/cli/aoti_*.py): torch.export and torch._inductor operate on raw tensors and pytrees by contract. Metadata is persisted at export and re-attached on load.
• torch.autograd.Function (e.g. _ImplicitUpdateFn): PyTorch's autograd machinery requires raw at the forward/backward signatures. The wrap-back happens immediately on exit from the enclosing typed method.
• Eager embed bridge (neml2/eager/, with the shared boundary helpers in neml2/types/_boundary.py): the C++ neml2::eager::Model marshals raw at::Tensors across the embedded-Python boundary via torch's pybind casters. Raw tensors appear only in forward / jvp / jacobian's argument / return dicts; inputs are wrapped to typed immediately on entry and outputs are unwrapped only at the return (in the neml2/types/_boundary.py unwrap_outputs / assemble_jvp_outputs / assemble_jacobian helpers).
• pyzag adapter (neml2/pyzag/): the pyzag time-integration library's BlockVector / SolvableBlockOperator / BlockJacobian / NonlinearFunctionOperatorFactory ABCs operate on raw torch.Tensors by contract. Raw appears only where one genuinely crosses that ABC surface — constructing/consuming BlockVector.raw_tensors, feeding torch.linalg, or returning to a pyzag ABC method. Reading metadata or manipulating neml2-typed objects stays on the typed API (add an ES / types primitive when one is missing); the assembled residual/Jacobian is produced typed through neml2.es and only unwrapped at the final pyzag hand-off.

Everywhere else, returning raw is a bug.

2. No .data access outside neml2/types/

The .data attribute exists for the type implementations themselves. Every other caller is forbidden from touching it. If you need an op that the typed wrappers don't expose, add the op inside neml2/types/ and use it from your call site — don't reach around the typed API.

A .data access elsewhere means one of:

• The caller is about to do raw torch.* arithmetic and rewrap, dropping labels/state/meta and reconstructing them by guess. This is the root cause of metadata-loss bugs.
• A wrapper primitive is missing. Add it.

The same framework boundaries (AOTI shim, torch.autograd.Function, eager embed bridge) are the only exceptions, and even there the unwrap is encapsulated in a helper inside neml2/types/ (the shared boundary helpers live in neml2/types/_boundary.py), not done ad hoc at the call site. Genuine framework-boundary unwraps in those exception files bear a # data-ok marker explaining why.

The pyzag adapter (neml2/pyzag/) is an additional boundary of this kind, but — like the existing in-tree boundaries neml2/es/implicit.py, neml2/models/param_ad.py, and neml2/cli/aoti_export.py — it accesses .data directly at the call site with an x.data # data-ok pyzag boundary marker rather than through a shared helper. This is sanctioned only where a raw tensor actually crosses into pyzag's ABC surface (constructing/consuming BlockVector.raw_tensors, feeding torch.linalg, returning to a pyzag ABC method). It is not sanctioned for reading metadata (.device / .dtype / .shape / .batch_shape — use the typed accessors) or for manipulating neml2-typed objects (transpose, indexing, flip/clone/pad, matvec — use the typed AssembledMatrix / AssembledVector / Tensor API, adding a primitive if one is missing).

3. Fix the root cause; never patch the call site

When you find a bug, trace upstream until you reach the source of the wrong behaviour, not the place where it became visible. Patching the consumer warrants patches at every other consumer hitting the same root, and the codebase ends up with N scattered workarounds for one underlying bug. If the answer is "I'll patch this site for now and fix upstream later", choose fix upstream now — "later" almost never happens and the patch becomes load-bearing.

A useful test: would a hypothetical new caller hit this same bug? If yes, the bug is in what they're calling. Fix it there.

These three rules reinforce each other. Rule 1 forces functions to produce typed outputs. Rule 2 forces typed operations to exist for every need. Rule 3 forces the fix to land at the wrapper/producer level once, instead of N times at each consumer.

Build & develop

pip install torch nmhit scikit-build-core cmake ninja   # build prerequisites (see note)
pip install -e ".[dev]" -v --no-build-isolation

This drives a scikit-build-core build of the C++ runtime under build/<wheel_tag>/ (stable build/editable symlink), then installs the Python package in editable mode. Python source edits take effect immediately; touching anything under neml2/csrc/ or CMakeLists.txt needs a re-pip install — or cmake --build build/editable for a fast incremental C++ rebuild. There are no CMake presets: pip install -e is the single build entry point.

Why the manual prerequisites + --no-build-isolation. torch and nmhit are declared as runtime deps in [project.dependencies], but pip installs runtime deps after it builds the wheel — and the C++ build needs libtorch + libnmhit (CMake find_package) during the build. They can't move to build-system.requires either: libneml2.so's rpath resolves libtorch from the sibling torch/ package at runtime, so neml2 must compile against the same torch you run, not an isolated build-env copy. So the build has to see your environment's torch/nmhit — hence pre-installing them and --no-build-isolation (which, in turn, requires the build backend — scikit-build-core + cmake + ninja — to be present too). This mirrors cibuildwheel (build-frontend = build --no-isolation + a before-build install).

An editable install builds at RelWithDebInfo and additionally compiles the C++ test executables (the cpp_tests aggregate target) and emits compile_commands.json (symlinked to the repo root); a shipped wheel (SKBUILD_STATE != editable) stays lean (Release, no tests). The editable-vs-wheel split is declared via [[tool.scikit-build.overrides]] if.state = "editable" in pyproject.toml.

Package versions and pinned deps live in dependencies.yaml — use python scripts/dep_manager.py {check|list|bump DEP.FIELD VALUE} rather than editing version strings by hand. Files reference their dep with a # dependencies: NAME.FIELD annotation immediately above the version literal. The torch compatibility matrix (compatibility.yaml) is a separate registry checked against the same dependencies.yaml torch entry — keep them in sync via python scripts/compat_matrix.py {seed|check|render}. Never hardcode a dependency version in prose or documentation (no "CMake ≥ 3.26", no "torch 2.12") — reference the dependency by name; the pinned value lives only in dependencies.yaml and the annotated literals, so a bump never has to chase prose.

C++ tests + instrumented builds

The C++ test executables build with the editable install; run them by ctest label against build/editable:

ctest --test-dir build/editable -L dispatcher    # scheduler / dispatcher tests
ctest --test-dir build/editable -L eager         # embedded-Python eager test
ctest --test-dir build/editable -L benchmark     # benchmark smoke tests

The instrumented build types go through the same entry point via --config-settings:

# Coverage (clang source-based) -> coverage.lcov
CC=clang CXX=clang++ pip install -e ".[dev]" --no-build-isolation \
  --config-settings=cmake.build-type=Coverage
scripts/cpp_coverage.sh build/editable                       # runs all ctests except the `benchmark` label

# ThreadSanitizer (guards the async dispatch pool)
CC=clang CXX=clang++ pip install -e ".[dev]" --no-build-isolation \
  --config-settings=cmake.build-type=ThreadSanitizer
OMP_NUM_THREADS=1 TSAN_OPTIONS="suppressions=tests/cpp/tsan_suppressions.txt ignore_noninstrumented_modules=1" \
  ctest --test-dir build/editable -L dispatcher

The pip libtorch is not TSan-instrumented, so ignore_noninstrumented_modules=1 + tests/cpp/tsan_suppressions.txt silence torch's own reports; neml2's code is still checked. CI runs both on ubuntu+clang (the cpp-coverage and tsan jobs in .github/workflows/cpp.yaml); C++ coverage uploads to Codecov under the informational cpp flag.

Tests

All tests are pytest. Test layout under tests/ is fixed at five top-level buckets:

• tests/unit/ -- Python unit tests of individual modules (typed wrappers, chain rule, factory, schema, solvers, drivers, AOTI export metadata, CLI extensions, pyzag adapter). Fast.
• tests/models/ -- ModelUnitTest-driven .i files exercising one registered model leaf at a time. Discovered by tests/models/test_model_unit_tests.py.
• tests/regression/ -- parametrized regression sweep that pins each scenario's output against a checked-in gold/result.pt reference.
• tests/verification/ -- parametrized verification sweep that compares scenario output against an external ground truth.
• tests/aoti/ -- AOTI end-to-end smoke tests (compile, load, run one scenario per pattern). Slow -- each test triggers an Inductor compile -- but runs by default.

pytest tests/                                  # everything
pytest tests/unit/                             # the fast unit suite
pytest tests/unit/test_factory.py              # one file
pytest tests/unit/test_factory.py::test_load_input_simple   # one function
pytest tests/regression/                       # the parametrized regression sweep
pytest tests/verification/                     # the parametrized verification sweep
pytest tests/aoti/                             # AOTI compile smoke suite
pytest -n auto tests/aoti/                     # AOTI suite parallel (4.6x speedup at -n 8)
pytest --cov tests/unit tests/models tests/aoti   # branch+line coverage report

tests/regression/test_regression.py and tests/verification/test_verification.py discover scenarios by walking their respective directories for .i files and emit one parametrize id per scenario (the input file's relative path). For spot-checks after a code change, target a single scenario via -k or the full parametrize id:

pytest tests/regression/test_regression.py -k maxwell
pytest 'tests/regression/test_regression.py::test_regression[solid_mechanics/viscoelasticity/maxwell/model.i]'

Reserve unfiltered tests/regression/ / tests/verification/ runs for final confirmation.

When adding a Model subclass, use the /add-model skill; for a regression or verification scenario use /add-regression or /add-verification. The skills encode the test conventions so you don't have to rediscover them.

Documentation

Sphinx with the shibuya theme and MyST-NB. Use the wrapper script:

doc/scripts/build.sh            # parallel build, -W, --keep-going, html → doc/_build/html
doc/scripts/build.sh --clean    # wipe doc/_build first (cold build)
doc/scripts/build.sh --serve    # build + serve at http://127.0.0.1:8765/ (binds 127.0.0.1)

The script wraps sphinx-build -j auto -W --keep-going and pre-creates doc/_build/html/.jupyter_cache so myst-nb's parallel workers don't race the directory creation. Pass --port N, -j N, --no-strict, or --dest PATH to override defaults; --help for the full list.

The /build-docs skill walks through the full pipeline (including notebook execution caching and the auto-generated HIT-syntax catalog). The neml2 package must be importable for autodoc / neml2-syntax to introspect the registered objects — an editable pip install -e ".[dev]" is enough.

CLI tools

The installed wheel exposes four console scripts (defined in pyproject.toml under [project.scripts]):

• neml2-run <input.i> — drive a model through a load history.
• neml2-inspect <input.i> — print the resolved input/output graph of a wired-up input file. Use this before neml2-run when composing models; wiring bugs surface as obvious mismatches instead of cryptic shape errors deep in Newton.
• neml2-syntax --section Models --summary — browse the registered-object catalog with one-line docstrings (--type <Name> to drill into one). Run this when planning any new Model or wondering whether a primitive already does what you want.
• neml2-compile <input.i> --model <name> — export a model to an AOT-Inductor .pt2 package + drop-in HIT stub. See for the artifact format and for the compilation pipeline.

neml2-diagnose and neml2-time from v2 are gone.

Architecture

The Python package layout under neml2/:

• factory.py — HIT input parsing via nmhit, load_input / load_model / load_nonlinear_system entry points, [Tensors] namespace for inline Python expressions.
• schema.py — declarative HIT syntax: input, output, parameter, option field helpers + HitSchema container that drives both parsing and the auto-generated docs.
• models/ — the composable forward operators and the framework infrastructure that powers them:
  • models/model.py — Model base class (a torch.nn.Module). All user-authored constitutive leaves inherit from this.
  • models/chain_rule.py — type aliases for the first / second-order chain-rule sensitivity dicts threaded through Model.forward(..., v=, v2=, vh=).
  • models/param_ad.py — reverse-mode parameter derivatives (param_jacobian / param_vjp via torch.autograd.grad), the reverse-mode counterpart to chain_rule.py's forward-mode input sensitivities. The single engine shared by every route (native Model.param_jacobian/param_vjp, _EagerModel, and — wrapped for export — the AOTI graphs).
  • models/resolver.py — DependencyResolver builds the ComposedModel dependency graph from individual leaves' declared inputs and outputs.
  • models/export.py — adapter around torch.export + torch._inductor.aoti_compile_and_package; the single entry point through which every AOTI lowering passes.
  • models/_guard.py — forward() guard that blocks raw torch.autograd / einsum calls inside leaves; the allow_autograd / allow_einsum context managers carve out exceptions.
  • models/common/ — ComposedModel glues children together via the dependency graph; ImplicitUpdate wraps a residual model in a Newton solve with optional Predictor.
  • models/{solid_mechanics,chemical_reactions,phase_field_fracture,porous_flow,finite_volume,kwn}/ — domain leaf libraries. Crystal plasticity is a subdirectory of solid_mechanics/.
• types/ — typed tensor wrappers (Scalar, Vec, R2, SR2, MRP, Quaternion, MillerIndex, fourth-order SSR4 / WSR4 / ...). Each is a dataclass registered with torch.utils._pytree.register_dataclass so it round-trips through torch.export. .data exposes the underlying torch.Tensor. types/_boundary.py holds the shared raw-tensor framework-boundary helpers (device/dtype check_tensor, broadcast_to_common_batch, and the typed-from-raw unwrap_outputs / assemble_jvp_outputs / assemble_jacobian) used by the AOTI shim, the eager bridge, and Model — not eager-specific, so it lives here at the raw↔typed boundary rather than in any one consumer.
• solvers/ — package with dense_lu.py, schur_complement.py, newton.py, newton_linesearch.py (per-class files mirroring v2's layout).
• es/ — equation-systems package with axis_layout.py, assembled.py (AssembledVector/Matrix wrapping the dynamic-base Tensor), system.py (LinearSystem / NonlinearSystem / ModelNonlinearSystem), and implicit.py (AOTI implicit-segment export wrappers RHS / NewtonStep / IFT). The same three wrappers cover the DenseLU and SchurComplement solver paths -- the linear solver is configured externally and forwarded.
• drivers/ — driver.py (the abstract Driver base) plus the concrete TransientDriver, ModelUnitTest, TransientRegression, Verification files — the top-level "run a model over a load history" objects exposed in input files.
• data/ — CubicCrystal, CrystalGeometry and related crystallography data classes.
• user_tensors/ — registered [Tensors] block types other than Python (currently the CSV<Type> family).
• cli/aoti_compile.py, cli/aoti_export.py — the neml2-compile orchestration and the per-segment export path; see .
• aoti/ — Python-side AOTIModel shim that loads the shared metadata.json + per-<device>/<dtype>/ .pt2 binaries produced by neml2-compile and exposes forward / jvp / jacobian. Backed by the pybind module aoti/_aoti.cpp (which links libneml2.so).
• eager/ — the Python adapter package the C++ embedded-Python eager runtime imports (the import path neml2.eager is stable; __init__.py re-exports _EagerModel from eager/_model.py). _EagerModel wraps a factory.load_model native model and presents the raw-tensor, name-keyed forward / jvp / jacobian (dict→dict, plus (dict, J) for jacobian) + input_names/sizes/device/dtype surface the C++ neml2::eager::Model consumes; jvp/jacobian reuse the native model's v= chain rule (the _ForwardJacobianModule seed/assembly helpers). The shared device/dtype check + batch-broadcast + output-unwrap/assembly helpers live in neml2/types/_boundary.py (the third raw-tensor framework boundary; see Rule 1) — they are not eager-specific (the AOTI shim and Model use them too). Eager is plain-batch only: a model that produces sub-batch output (e.g. crystal plasticity) is rejected (no slot to declare per-input sub-batch shapes at this boundary).
• pyzag/ — adapter exposing a NEML2 nonlinear system to the pyzag time-integration/training library. interface.py's NEML2PyzagModel wraps a neml2.es.ModelNonlinearSystem as a pyzag.nonlinear.NonlinearFunctionOperatorFactory (assembling the per-chunk residual + bidiagonal Jacobian, mirroring HIT parameters as torch.nn.Parameters). pyzag/operators/ implements pyzag's BlockVector/SolvableBlockOperator/BlockJacobian on top of AssembledVector/AssembledMatrix — cached-LU Thomas + SchurComplement solves; parallel cyclic reduction is supported for single-group DENSE layouts (via pyzag's dense backend), while BLOCK (per-site, e.g. crystal-plasticity) layouts use the Thomas factorization.
• cli/ — backing modules for the four console scripts above.
• csrc/aoti/ — C++ runtime: neml2::aoti::Model wraps torch::inductor::AOTIModelPackageLoader. The public class is a PImpl facade (Model.h); its internals live behind Model::Impl in the non-shipped internal.h (+ assertions.h), and the implementation is split across Model.cpp (construction), ops.cpp (forward/jvp/jacobian), solve.cpp (value/Newton path), jacobian.cpp (Jacobian/IFT path), and the shared newton.{h,cpp} / nonlinear_system*.{h,cpp} solver. Exception.h (shipped) is the public exception taxonomy — Exception base with a recoverable() predicate, ConvergenceError (recoverable: a Newton divergence / max-iters, so a consumer can cut the time step and retry), FatalError (non-recoverable: shape/device/config; what _assert throws), and AggregateError (concurrent dispatch failures). Public ops run through _guarded so foreign torch errors are normalized to FatalError. Built into neml2/lib/libneml2.so (hidden visibility; only the AOTI_EXPORT-tagged API is exported) and surfaced through the pybind binding.
• csrc/dispatchers/ — C++ work scheduler / dispatcher (serves the compiled path embedded in a host app; no Python). DispatchedModel is a Model-shaped handle owning one pinned aoti::Model per device + an injected WorkScheduler; it chunks a batched call across devices and stitches results back. Schedulers split into SyncScheduler (SimpleScheduler, MPISimpleScheduler — single-device chunk loop on the calling thread) and AsyncScheduler (StaticHybridScheduler — concurrent CPU+GPU(s) via a thread-per-device pool, load-tracked). factory.cpp is the load_model(stub, name[, scheduler]) entry point; batch_chunk.h holds the slice/cat helpers. See .
• csrc/eager/ — C++ embedded-Python eager runtime, compiled into the separate neml2/lib/libneml2_eager.so (the only NEML2 C++ artifact that links torch_python; libneml2.so stays Python-free). It links Python3::Module (not libpython) and leaves the CPython API symbols undefined — resolved at load from the host's libpython (the running interpreter, or the libpython a pure-C++ host links to embed it); this keeps the wheel build off any shared libpython dependency. neml2::eager::Model is a PImpl facade (Model.h) mirroring aoti::Model's forward / jvp / jacobian + metadata surface, but constructed from the original .i (no compile): it embeds a CPython interpreter (one-time bootstrap in interpreter.{h,cpp}, never finalized — torch can't be re-imported after Py_Finalize), imports the neml2.eager._EagerModel Python adapter, and marshals at::Tensors across the boundary with torch's pybind casters (Model.cpp). load_model(input_file, name) (load_model.cpp) is the entry point. Python failures are normalized via guarded() to neml2::aoti::FatalError, except a solver divergence / max-iters, which round-trips as the recoverable neml2::aoti::ConvergenceError (it surfaces from libneml2.so as the neml2.aoti._aoti.ConvergenceError registered in _aoti.cpp). Plain-batch only (sub-batch models are rejected). The fast-to-start, slow-to-run counterpart of the AOTI path, for downstream C++ unit tests. See .

Factory / Registry pattern

Every concrete native object (model, driver, tensor, solver, ...) self-registers via the @register_neml2_object("TypeName") decorator from neml2.factory. HIT input files then instantiate them by type name. Model subclasses declare their input/output/parameter surface via a class-level hit = HitSchema(...) and a from_hit constructor; the @register_neml2_object decorator + HitSchema together feed both the live factory and the auto-generated neml2-syntax catalog.

When adding a new submodule under neml2/models/<domain>/, append the import to the parent __init__.py so import neml2 triggers registration — there is no lazy-loading machinery and unimported modules' types are invisible to the factory.

Conventions

• Python source: linted and formatted with ruff (line length 100), CI-enforced via the lint job in .github/workflows/python.yaml. Run pre-commit run --all-files before pushing.
• Type-checked with pyright against the installed package (CI: the typecheck job).
• Math in docstrings uses MyST dollarmath ($x$ inline, $$...$$ display); MyST dollarmath and amsmath extensions are enabled in doc/conf.py. Code references stay in `backticks`; the difference matters for rendering in the syntax catalog.
• HIT inputs use the nmhit Python parser (also a pre-commit nmhit-format hook). The format itself is unchanged from v2.
• Tutorials under doc/content/tutorials/**/main.ipynb are notebook-only (no jupytext pairing — edit the .ipynb directly; review via GitHub's notebook diff). Cheap tutorials are executed at build time (nb_execution_mode = "cache") and committed without outputs; the two expensive pyzag notebooks (optimization/{deterministic,statistical}) are committed pre-baked, excluded from build-time execution, and kept current by the check-notebook-executed hook. Reference tutorial pages with no .ipynb (e.g. index.md, models/input_file.md) are plain markdown.
• Copyright headers are checked by a pre-commit hook (python scripts/check_copyright.py); the script auto-fixes missing headers.
• Avoid editing files in build/, installed/, or doc/_build/, doc/generated/ — those are generated. scripts/clobber.sh [dir] removes git-ignored files if a build gets wedged.