neml2.cli

Native CLI entry points for the neml2-* tools.

One module per tool, each exposing main(argv=None) -> int. Tools are registered as console scripts in pyproject.toml’s [project.scripts].

export_model_for_aoti is also re-exported here so the public symbol neml2.export_model_for_aoti keeps resolving after the aoti_export module moved into this subpackage.

neml2.cli.export_model_for_aoti(hit_path, model_name, output_dir, *, device='cpu', dtype='float64', promoted=(), example_batch_shape=None, dynamic_batch=None, derivatives=(), pre=(), additional_args=())[source]

Export a native NEML2 model to .pt2 artifacts for AOTI C++ consumption.

By default every parameter and buffer in the source model is folded into the exported graph as a constant – the most efficient form for inference use. To keep a parameter mutable at runtime, pass its fully-qualified name via promoted; the exporter traces those entries as additional graph inputs and records their initial values in the metadata so the C++ side can populate aoti::Model::named_parameters().

Parameters:

  • hit_path (str | Path) – Path to the HIT .i file.

  • model_name (str) – Name of the model in the [Models] section.

  • output_dir (str | Path) – Directory for .pt2 artifacts and metadata JSON. Created if absent.

  • device (str) – Target device for the resulting artifact (default "cpu").

  • dtype (str) – Floating-point dtype for the resulting artifact ("float64" or "float32"; default "float64").

  • promoted (set[str] | list[str] | tuple[str, ...]) – Set / iterable of fully-qualified parameter or buffer names (the same form model.named_parameters(recurse=True) emits) to promote to runtime-flexible status. Empty default = fully baked artifact.

  • example_batch_shape (dict[str, str] | str | None) – Override for the per-input example shape used at trace time. Accepts either a single spec string (uniform across all inputs) or a dict mapping input variable name to spec string (per-variable). Each spec uses the (dyn; sub) grammar (see _parse_example_batch_shape()). Overrides any [Settings]/example_batch_shape in the HIT file. When None (default), the HIT setting – or _DEFAULT_EXAMPLE_SHAPE if unset – applies.

  • dynamic_batch (bool | None) – Override for [Settings]/dynamic_batch. True lets the leading batch axis vary at runtime (the default); False produces a static-batch artifact pinned at the example shape. Use False when a baked parameter has a rank ≥ 1 shape that would specialize the dynamic dim. When None (default), the HIT setting – or True if unset – applies.

  • derivatives (Sequence[str]) – Sequence of OUT:IN derivative specs (the -d/--derivative CLI surface). Each requests a Jacobian/JVP block for that output-input pair; omitting a side selects all on that side (stress: = all inputs of stress, :strain = all outputs wrt strain, : = all pairs). Empty (the default) compiles no derivative graphs, and the runtime jvp / jacobian raise until recompiled with -d. Resolved against the model’s outputs + structural inputs after promotion (see _resolve_derivative_specs()).

  • pre (Sequence[str]) – HIT snippets prepended before parsing (same semantics as nmhit.parse_file’s pre arg). Use to bind ${var} substitution variables that the input file expects – e.g. pre=["nbatch=2", "device=cpu"]. Post-:= overrides via additional_args do NOT propagate into ${...} substitution.

  • additional_args (tuple[str, ...]) – Trailing HIT override tokens appended to the parser’s post list (e.g. "Models/elasticity/E:=210000"). Path-style overrides only; for variable substitution use pre.

Returns:

The metadata dictionary (same content as the written JSON).

Return type:

dict