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 exposes forward, jvp, jacobian plus the parameter-derivative pair param_jacobian / param_vjp (d(output)/ d(parameter) for promoted parameters) – all 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. The derivative graphs are emitted only for the (output, input) / (output, parameter) pairs requested with neml2-compile --derivative OUT:IN (jvp / jacobian / param_jacobian / param_vjp raise for pairs that were not compiled in).

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) 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).

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, exposed via the 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>
jacobian(inputs, param_overrides=None)[source]

Outputs + the per-(output, input) Jacobian blocks from the compiled jacobian graph.

inputs is a {variable_name: tensor} dict. Returns (outputs, J) where J[out][in] is the block d(out)/d(in), shaped (batch, *out_base, *in_base). Like jvp(), needs an artifact compiled with -d.

Parameters:
Return type:

tuple[dict[str, Tensor], dict[str, dict[str, Tensor]]]

jvp(inputs, tangents, param_overrides=None)[source]

Outputs + the directional derivative J @ v from the compiled jvp graph.

inputs / tangents are {variable_name: tensor} dicts (same keys as input_spec); tangents carries the seed direction per input. Returns (outputs, jvp) keyed by output name – the same raw-tensor surface the eager runtime exposes. Requires the artifact to carry the derivative graph (compiled with neml2-compile ... -d OUT:IN); otherwise the underlying binding raises.

Parameters:
Return type:

tuple[dict[str, Tensor], dict[str, Tensor]]

named_parameters()[source]

The promoted parameters as a mutable {qualified_name: tensor} dict.

Overrides nn.Module.named_parameters – whose (name, Parameter) iterator would be empty here, since the shim registers no nn.Parameter``s of its own (the calibratable values live in the compiled binding). This returns the binding's promoted-parameter surface, the same dict the eager runtime and the C++ routes expose. Entries may be mutated in place (``m.named_parameters()["elasticity.E"].fill_(...)) or replaced via set_parameter(); the change is seen on the next call. Only parameters promoted at compile time (neml2-compile ... -p NAME) appear – with none promoted the dict is empty.

Return type:

dict[str, Tensor]

param_jacobian(inputs, param_overrides=None)[source]

Outputs + the per-(output, parameter) Jacobian blocks d(out)/d(param).

The reverse-mode parameter-derivative counterpart to jacobian() (mirrors neml2.models.model.Model.param_jacobian() on the other routes): inputs is a {variable_name: tensor} dict and the result is (outputs, P) with P[out][param] the block d(out)/d(param), keyed by qualified promoted-parameter name (see named_parameters()). Needs an artifact with the parameter promoted (neml2-compile ... -p NAME) and its derivative graph (-d).

Parameters:
Return type:

tuple[dict[str, Tensor], dict[str, dict[str, Tensor]]]

param_vjp(inputs, cotangents, param_overrides=None)[source]

Reverse-mode parameter VJP – cotangent contracted with d(out)/d(param).

The reverse-mode counterpart to param_jacobian() (mirrors neml2.models.model.Model.param_vjp()): inputs is a {variable_name: tensor} dict, cotangents is keyed by output name, and the result is keyed by qualified promoted-parameter name. Needs a -p / -d artifact.

Parameters:
Return type:

dict[str, Tensor]

set_parameter(name, value)[source]

Replace a promoted parameter’s value (forwarded to the binding).

name is a qualified promoted-parameter name (a key of named_parameters()); value must match the artifact’s pinned device/dtype. The new value is used on the next forward / derivative call.

Parameters:
Return type:

None

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], param_overrides: 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. param_overrides (default empty) replaces a promoted parameter’s value for this call only, without mutating named_parameters() – a hook for multi-device dispatch.

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], param_overrides: 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], param_overrides: 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.

param_jacobian(self: neml2.aoti._aoti.Model, inputs: collections.abc.Mapping[str, torch.Tensor], param_overrides: collections.abc.Mapping[str, torch.Tensor] = {}) → tuple[dict[str, torch.Tensor], dict[str, dict[str, torch.Tensor]]]

Evaluate + parameter Jacobian as unflattened variable-pair blocks (schema v7).

Returns a 2-tuple (outputs, P) where P is a nested dict[str, dict[str, Tensor]]: P[out_name][param_qname] is the block (*B, *out_base_shape, *param_base_shape) (e.g. stress w.r.t. a Scalar E -> (*B, 6)). The keys are promoted parameters (see named_parameters()), not structural inputs. Requires the artifact was compiled with -d OUT:PARAM over a promoted parameter; otherwise raises.

param_vjp(self: neml2.aoti._aoti.Model, inputs: collections.abc.Mapping[str, torch.Tensor], cotangents: collections.abc.Mapping[str, torch.Tensor], param_overrides: collections.abc.Mapping[str, torch.Tensor] = {}) → dict[str, torch.Tensor]

Parameter VJP / adjoint (schema v7): dL/d(param) for the loss L = sum_o <cotangent_o, out_o>.

cotangents is a dict[str, Tensor] keyed by output name, each at the output’s (*B, *out_base_shape) shape. Returns dict[str, Tensor] keyed by parameter qualified name – the cheaper form for many-parameter inverse optimization. Same compile requirement as param_jacobian.

property parameter_base_shapes

Per-promoted-parameter natural base shape, keyed by qualified name (Scalar -> [], SR2 -> [6]). The keys are the promoted parameters (same keys as named_parameters()); empty when nothing was promoted. The parameter analogue of input_base_shapes / output_base_shapes, and the unified parameter surface shared with neml2.eager._EagerModel.

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.