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 per-device artifact folder produced by neml2-compile (the folder holding one <device>/ subfolder per compiled device). The subfolder for the current torch.get_default_device() is loaded – so neml2-run --device cuda (which sets the default device) picks cuda/. 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

Absolute path to the per-device artifact folder produced by neml2-compile (contains one <device>/ subfolder per compiled device). The subfolder matching torch.get_default_device() is loaded.

solver — str

Solver whose convergence / line-search settings configure the implicit Newton solve. Schema v4+ no longer bakes these into the artifact; the stub .i carries the [Solvers] block and it is forwarded to the C++ runtime at load. Defaults apply for forward-only models.