AOTIModel

Source: aoti/_shim.py

HIT-loadable wrapper around neml2.aoti.Model.

Constructed from a HIT [Models] block with an artifact_path option pointing at the artifact root folder produced by neml2-compile (the folder holding one shared metadata.json and per-<device>/<dtype>/ .pt2 binaries). The <device>/<dtype>/ leaf is a load-time choice: load_model(..., device=, dtype=) (or neml2-run --device/--dtype) selects it, defaulting to NEML2’s canonical cpu / float64 (matching neml2-compile) rather than torch’s ambient defaults – so a stock artifact loads without the caller having to match torch’s float32 default. Eager and single-device: no dispatch happens here.

Plays the native-Model role: input_spec and output_spec are populated from the metadata’s var_type fields; __call__ takes TensorWrapper positional args in input_spec order, unwraps them to raw tensors, runs the AOTI forward graph, and wraps each output back in its declared type. Promoted parameters (if any) are not part of input_spec; they are reachable through named_parameters (a mutable dict) and set_parameter. The full sensitivity surface – jvp, jacobian, param_jacobian, param_vjp – is forwarded to the binding so the py-aoti route matches the others (CLAUDE.md “six evaluation routes” parity).

Other options

artifact_path — str · required

Path to the artifact folder produced by neml2-compile – one shared metadata.json plus <device>/<dtype>/ .pt2 binaries. Written by neml2-compile relative to the stub (so the stub + folder relocate together); resolved against the stub’s directory, and an absolute path is also accepted. The stub is device/dtype-agnostic (only artifact_path); the <device>/<dtype>/ leaf is chosen at load time via load_model(..., device=, dtype=), defaulting to cpu / float64 (matching neml2-compile). Solver config is read from the shared metadata, so no [Solvers] block is needed.