cpp-eager — eager evaluation from C++#

The compiled path (AOTI packages) is the right way to run a NEML2 model from C++ in production: it is fast and needs no Python at the hot loop. But neml2-compile runs an AOT-Inductor compile that takes minutes — too slow to sit in a downstream C++ project’s unit tests, where you just want to load a model and check a few values.

The eager C++ runtime fills that gap. neml2::eager::Model embeds a CPython interpreter, loads the model in Python via neml2.factory.load_model (the very same path neml2-run uses), and evaluates forward by marshalling at::Tensors across the C++/Python boundary. There is no compile step and no artifact — it consumes the original .i file directly. The trade-off is speed: every call goes through Python, so it is far slower than a compiled graph. Use it for fast-to-start tests; use the compiled / dispatched path (cpp-dispatch — dispatching across devices) for production.

Note

This runtime lives in a separate shared library, libneml2_eager.so. The main libneml2.so (the AOTI runtime + dispatcher) stays Python-free so it can embed in pure-C++ hosts; link libneml2_eager only where you want the eager fallback. libneml2_eager itself does not link libpython — like a CPython extension module, it leaves the Python C-API symbols undefined and resolves them at load time from the host’s interpreter. A pure-C++ host must therefore link libpython itself (e.g. CMake Python3::Python from find_package(Python3 COMPONENTS Development.Embed)); a host that already runs Python supplies it automatically.

Load and run from C++#

neml2::eager::load_model(input_file, model_name) mirrors Python’s load_model(path, name) and returns a neml2::eager::Model whose operations (forward / jvp / jacobian plus the parameter derivatives param_jacobian / param_vjp) and metadata accessors (input_names / output_names / input_base_shapes / output_base_shapes / input_sub_batch_shapes / output_sub_batch_shapes / device / dtype) have the same signatures as neml2::aoti::Model. So test code switches between the eager and compiled runtimes by changing only the load call and the header.

Listing 7 tests/aoti/forward_single/model.i — loaded directly, no compile#
# forward_single: a one-leaf forward-only model. Smoke test for the
# bake-by-default path on the simplest possible shape.

[Models]
  [model]
    type = LinearIsotropicElasticity
    strain = 'strain'
    stress = 'stress'
    coefficients = '100 0.3'
    coefficient_types = 'YOUNGS_MODULUS POISSONS_RATIO'
  []
[]

#include "neml2/csrc/eager/load_model.h"

using namespace neml2::eager;

// Loads the model named "model" straight from the .i by embedding Python.
auto m = load_model("tests/aoti/forward_single/model.i", "model");

// Build each input at its canonical (*B, *base_shape) shape, on the model's
// device + dtype (NEML2 runs in float64). For SR2 strain the base shape is {6}.
const auto opts = at::TensorOptions().dtype(m.dtype()).device(m.device());
std::map<std::string, at::Tensor> inputs;
for (std::size_t i = 0; i < m.input_names().size(); ++i)
{
  std::vector<int64_t> shape{8};                       // batch
  const auto & base = m.input_base_shapes()[i];        // e.g. {6} for SR2, {} for Scalar
  shape.insert(shape.end(), base.begin(), base.end());
  inputs.emplace(m.input_names()[i], at::randn(shape, opts));
}

auto out = m.forward(inputs);   // {"stress": (8, 6) tensor}

The reported input_names / input_base_shapes (and the outputs) are derived through the same helper the AOTI metadata uses, so an eager model and its compiled twin agree on the boundary contract — making neml2::eager::Model a true drop-in for neml2::aoti::Model. Inputs must be passed at their canonical (*B, *sub_batch_shape, *base_shape) shape; a non-canonical trailing shape (e.g. an SR2 strain passed as (*B, 1) instead of (*B, 6)) is rejected with a FatalError. Outputs come back at (*B, *out_sub_batch_shape, *out_base_shape), contiguous, as on neml2::aoti::Model.

Sub-batch#

The sub-batch extents come from the input file’s [Settings]/example_batch_shape (Sub-batch axes), the same declaration neml2-compile reads. The block stays optional:

  • declared: each declared input is split into (*B, *sub_batch, *base) with the declared extent, and a tensor that does not carry it is rejected (a size-1 sub-batch axis is broadcast to the extent). An unknown name in the declaration fails the load, as it fails neml2-compile.

  • not declared: the input is plain batch, as before.

input_sub_batch_shapes() and output_sub_batch_shapes() report these declared extents, so an empty entry means “not declared as sub-batched”, not “plain-batch”. The output extents are what neml2-compile would record for the same declaration, worked out at load without compiling or solving, so a host can size sub-batched outputs before the first call. With a declaration the load therefore also runs the compile plan for these extents, which costs one more model load plus a zero-input probe of each forward segment, and a failure there fails the load.

That is the one place the two routes differ in meaning rather than in signature. On neml2::aoti::Model the extents were recorded from what the compiled graph produces, so they are definite and an empty entry means plain batch (see cpp-aoti — compiled model from C++). Without a declaration an eager model can still return a sub-batched output, e.g. one built on a sub-batched parameter grid; a host that supports both runtimes treats an undeclared variable as plain and checks the shapes it receives on the first eager evaluation, stopping with a clear message if an output carries an undeclared sub-batch axis.

From Python, neml2.eager._EagerModel’s forward / jvp / jacobian also take an optional batch_ndim, the number of leading batch axes; every axis between those and the base shape is then sub-batch, with or without a declaration (a declared extent is still checked). The C++ signatures do not take it.

Loading external (Python-authored) models#

Because the eager runtime is an embedded interpreter, a C++ host can drive models that live outside the installed neml2 package — the embedded-eager counterpart of the --load flag on neml2-run / neml2-compile. Pass the external Python extensions (file paths, package directories, or importable dotted module names) as the optional load argument; each is imported — registering its @register_neml2_object types into the factory — before the model is built:

// my_model.py defines `@register_neml2_object("MyModel") class MyModel(Model): ...`
// and my_model.i has `type = MyModel`. Without the load argument, MyModel is an
// unknown type and the load fails.
auto m = load_model("my_model.i", "model", {"path/to/my_model.py"});

Resolution and ordering match the CLI (neml2.cli._extensions.load_user_extensions). Registration is process-global: once an extension is loaded, its types stay visible to later load_model calls in the same process.

Sensitivities: jvp and jacobian#

jvp and jacobian mirror neml2::aoti::Model and are computed on the native model’s chain rule (the same v= machinery the AOTI export traces), so they work for forward and implicit (Newton) models alike. Both are variable-native: the boundary works in variables, not flattened group vectors.

// Jacobian-vector product: jvp_out[name] is the directional derivative at the
// output's natural (*B, *out_sub_batch, *out_base) shape.
auto [out, jvp_out] = m.jvp(inputs, tangents);   // tangents keyed + shaped like
                                                 // inputs; a missing key is zero.

// Full Jacobian as unflattened variable-pair blocks: J[out_name][in_name] is
// (*B, *out_sub_batch, *out_base, *in_base) -- e.g. SR2->SR2 -> (*B, 6, 6);
// Scalar->SR2 -> (*B, 6).
auto [out2, J] = m.jacobian(inputs);

Note

Derivatives are with respect to plain-batch inputs. A sub-batched input (e.g. per-bin history number_density~1) is held fixed: jacobian returns no column for it, as an AOTI artifact returns only the pairs it was compiled with, and jvp rejects a tangent for it. Rows for sub-batched outputs carry the output’s sub-batch axes. The parameter derivatives (param_jacobian / param_vjp) are plain-batch only and reject a call that carries a sub-batched variable.

Errors#

A Python-side failure — a parse error, an unknown model name, a wrong-dtype input, an input missing its declared sub-batch extent — is normalized to neml2::aoti::FatalError (the same taxonomy the AOTI runtime uses; see AOTI packages). The one exception is a solver divergence / max-iterations, which is re-raised as the recoverable neml2::aoti::ConvergenceError (it originates as that type inside libneml2.so and round-trips faithfully through the embedded interpreter), so a host can cut the time step and retry:

try { auto out = m.forward(inputs); }
catch (const neml2::aoti::Exception & e) {
  if (e.recoverable()) { /* cut dt, retry */ }
  else throw;
}

The exception types are exported from libneml2.so, so a catch (const neml2::aoti::Exception &) works across the two libraries.

Requirements#

Because the model is loaded in Python, the host process must be able to import the neml2 package and its dependencies:

  • A usable CPython interpreter. neml2::eager::Model starts one on first use (and never finalizes it — PyTorch does not support a clean re-import after Py_Finalize). If the host is already a Python process, the existing interpreter is reused and its lifecycle is left untouched.

  • The neml2 package importable from that interpreter (e.g. on PYTHONPATH, or installed in the interpreter’s environment).

Each call holds the GIL, so concurrent forward calls from multiple host threads are serialized. That is fine for unit tests; for parallel throughput use the compiled / dispatched runtime (cpp-dispatch — dispatching across devices).