neml2.aoti¶

Python interface to the thin C++ neml2::aoti::Model runtime.

This module exposes Model – a wrapper around the bare C++ class that loads AOTI-exported NEML2 model artifacts (.pt2 + _meta.json) produced by neml2-compile.

The runtime has three operations – forward, jvp, jacobian – keyed by the structural input/output names recorded in the metadata. A named_parameters() dict exposes parameters that were explicitly promoted via neml2-compile --parameter NAME at compile time; mutating those tensors in place is reflected on the next call.

Example usage:

from neml2.aoti import Model

m = Model("aoti/my_model/my_model_meta.json")
outputs = m.forward({"strain": strain_tensor})
outputs, J = m.jacobian({"strain": strain_tensor})

# If `--parameter E` was passed at compile time:
m.named_parameters()["E"].fill_(210000.0)

class neml2.aoti.AOTIModel(meta_path)[source]¶

Bases: Module

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.

Parameters:: meta_path (str | Path)

SECTION = 'Models'¶: Inherits from nn.Module rather than neml2.model.Model so the bound torch::inductor::AOTIModelPackageLoader runtime drives evaluation; the explicit class attribute keeps it in the [Models] section of the syntax catalog despite not subclassing Model.

forward(*args, v=None, v2=None, vh=None)[source]¶

Drive the AOTI graph.

Accepts input_spec positional args as TensorWrapper instances (mirroring the native ComposedModel boundary). Returns a tuple of typed wrappers in output_spec order – always a tuple, even for single-output models, so consumers can iterate uniformly.

v / v2 / vh (the native chain-rule hooks) are accepted only for signature compatibility with other native models; passing them is rejected because the AOTI graph’s JVP path is structurally different (it’s a separate _jvp.pt2 graph, accessed via the binding’s jvp() / jacobian() methods rather than the typed chain-rule protocol). Drivers that don’t need sensitivities – e.g. TransientDriver, TransientRegression – work as-is.

Parameters:: args (TensorWrapper)

classmethod from_hit(node, factory)[source]¶

Parameters:

node (nmhit.Node)
factory (_NativeInputFile)

Return type:

AOTIModel

hit = <neml2.schema.HitSchema object>¶

exception neml2.aoti.ConvergenceError¶: Bases: RuntimeError

class neml2.aoti.Model¶

Bases: pybind11_object

Thin C++ runtime for an AOTI-exported NEML2 model.

Construct from the path to the metadata JSON produced by neml2-compile; the loader resolves the per-segment .pt2 files relative to that path.

The artifact is device- and dtype-pinned at export time; there is no runtime to(). To target a different device, re-run neml2-compile.

Parameters that were explicitly promoted via --parameter NAME at compile time are reachable through named_parameters() and may be mutated in-place (e.g. model.named_parameters()['E'].fill_(210000.0)). Everything else is baked into the graph as a constant.

property device¶: Device the artifact was compiled for (immutable).

property dtype¶: Floating-point dtype the artifact was compiled for (immutable).

forward(self: neml2.aoti._aoti.Model, inputs: collections.abc.Mapping[str, torch.Tensor]) → dict¶

Evaluate the model.

inputs is keyed by the names returned by input_names; missing keys raise an error. Returns one tensor per name in output_names, preserving declaration order.

property input_base_shapes¶: Per-input base shape (Scalar -> [], SR2 -> [6], R2 -> [3, 3]). Inputs must be passed at their canonical (*B, *base_shape) shape.

property input_names¶: Master input names in graph-call order.

jacobian(self: neml2.aoti._aoti.Model, inputs: collections.abc.Mapping[str, torch.Tensor]) → tuple[dict[str, torch.Tensor], dict[str, dict[str, torch.Tensor]]]¶

Evaluate + full Jacobian as unflattened variable-pair blocks.

Returns a 2-tuple (outputs, J) where J is a nested dict[str, dict[str, Tensor]]: J[out_name][in_name] is the block (*B, *out_base_shape, *in_base_shape) (e.g. SR2->SR2 -> (B, 6, 6); Scalar->SR2 -> (*B, 6)) over the **structural* inputs (promoted-parameter inputs are not exposed in J).

jvp(self: neml2.aoti._aoti.Model, inputs: collections.abc.Mapping[str, torch.Tensor], tangents: collections.abc.Mapping[str, torch.Tensor]) → tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]]¶

Evaluate + JVP.

inputs and tangents are keyed by input_names and shaped (*B, *base_shape); a missing tangent key defaults to zero. Returns a 2-tuple (outputs, jvp_outputs) – both dict[str, Tensor] keyed by output_names; jvp_outputs[name] is the directional derivative at the output’s natural (*B, *out_base_shape).

named_parameters(self: neml2.aoti._aoti.Model) → dict[str, torch.Tensor]¶

Return the mutable map of runtime-flexible (promoted) parameters.

The dict’s tensor values share storage with the C++-side parameter slots; in-place mutation (model.named_parameters()['E'].fill_(...)) is reflected on the next forward / jvp / jacobian call. Reassigning a dict entry via model.named_parameters()['E'] = new_tensor updates the Python dict only, not the C++ slot – use set_parameter for that.

Empty when the model was compiled with no --parameter flags.

property output_base_shapes¶: Per-output base shape (Scalar -> [], SR2 -> [6], R2 -> [3, 3]).

property output_names¶: Master output names in graph-call order.

set_parameter(self: neml2.aoti._aoti.Model, name: str, value: torch.Tensor) → None¶: Replace a promoted parameter’s tensor (the C++-side slot is updated).

set_solver_config(self: neml2.aoti._aoti.Model, atol: SupportsFloat, rtol: SupportsFloat, miters: SupportsInt, ls_type: str, ls_max_iters: SupportsInt, ls_cutback: SupportsFloat, ls_c: SupportsFloat) → None¶: Configure the implicit-segment Newton solve (from the stub’s [Solvers] block). Schema v4+ no longer bakes these into the artifact.