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) live on the underlying binding’s named_parameters() and are not part of input_spec – the caller mutates them in place via self._inner.named_parameters() to drive runtime-flexible behavior.

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.