neml2.models.export

Export pipeline: torch.export → torch._inductor.aoti_compile_and_package.

This module wraps the leading-underscore torch._inductor APIs behind a single compile_model entry point so a torch version bump touches one file. The resulting .pt2 package is loadable from C++ via torch::inductor::AOTIModelPackageLoader (wrapped by neml2::aoti::Model) and from Python via load_package.

Pinned to torch 2.12.0 — see RISK.md R-001 and decision D-001.

neml2.models.export.compile_model(model, example_inputs, package_path, *, dynamic_batch_dim=0, batch_max=1048576, strict=False, inductor_configs=None)[source]

Export model and AOT-compile to a .pt2 package at package_path.

Parameters:

  • model (Module) – A torch.nn.Module whose forward takes/returns ``torch.Tensor``s (or tuples thereof). Buffers are baked in as constants; ``nn.Parameter``s are not supported by AOTInductor’s forward-only pipeline (see RISK.md R-007).

  • example_inputs (tuple[Any, ...]) – Positional tensors with representative dtype/device and a representative batch shape (typically batch=2 to force a true dynamic dim).

  • package_path (str | Path) – Destination path for the .pt2 artifact. Created if absent; overwritten if present.

  • dynamic_batch_dim (int | None) – Which leading axis of every input is dynamic (typically 0). If None, all input shapes are pinned and the export specializes.

  • batch_max (int) – Upper bound for the dynamic batch dim. PyTorch Dim requires explicit min/max; an unbounded max falls back to constant shape on some kernels. See RISK.md R-004.

  • strict (bool) – Passed to torch.export.export(). The default remains False for some models that rely on the more permissive tracer. Dense equation-system wrappers can use strict=True to avoid non-strict fake-tensor constant lifting around internally seeded tangent blocks.

  • inductor_configs (dict[str, Any] | None) – Optional dict forwarded to aoti_compile_and_package as inductor_configs=.... Use this to enable knobs like {"triton.cudagraphs": True} for CUDA-graph capture on graphs whose per-call kernel-launch overhead matters (e.g. DenseNewtonStep). None (the default) leaves Inductor at its compile-time defaults.

Returns:

The absolute path to the produced .pt2 artifact.

Return type:

Path

neml2.models.export.load_package(package_path)[source]

Load a .pt2 package produced by compile_model() for in-Python use.

This is the round-trip companion to compile_model(); the C++ side loads the same artifact via neml2::aoti::Model (include/neml2/aoti/Model.h).

Parameters:

package_path (str | Path)

Return type:

Any