NEML2
CLI utilities¶
The command line is the third way to use neml2 — alongside the Python (Python integration) and C++ (C++ integration) runtimes, you can drive a model end to end from a shell, with no code. The NEML2 wheel installs five console scripts:
Tool |
Purpose |
|---|---|
|
Drive a model through a load history from the shell. |
|
Print the structural summary of a model. |
|
Browse the registered-object catalog. |
|
Export a model to an AOT-Inductor package. |
|
Regenerate |
The first four share a common style — each takes an input file (where one is
needed), each forwards trailing positional tokens as HIT overrides
(Models/elasticity/E:=210000), each is implemented in pure Python
under neml2/cli/
and is also importable as a regular Python function.
All four of those tools accept a cumulative --load PATH flag for importing
user-defined extensions (custom Model / Driver / Solver /
[Tensors] / [Data] classes) before the tool’s own work begins.
PATH is either a path to a .py file or a package directory, or a
dotted module name on sys.path. The matching @register_neml2_object
decorators fire on import and the new types become resolvable from the
input file and visible to neml2-syntax. Repeat --load for several
extensions — they import in the order given, so later modules may
depend on names registered by earlier ones.
# A custom model defined in ./my_models.py is now driveable by neml2-run
# (and inspectable by neml2-inspect / neml2-compile, listed by neml2-syntax).
neml2-run --load ./my_models.py input.i my_driver
# Multiple extensions in priority order — later may depend on earlier:
neml2-syntax --load ./shared.py --load ./project_models.py --summary --json -
See Connecting to input files for the author’s side of the contract.
Most of the worked examples below reuse a single input file — a linear
isotropic elastic model named elasticity, wrapped in a five-step
TransientDriver named driver:
# Small worked example for the CLI utilities tutorial: a linear isotropic
# elastic model driven through a 5-step uniaxial tension history. The same
# input is reused across the neml2-run / neml2-inspect / neml2-compile
# subsections.
[Tensors]
[times]
type = Python
expr = 'Scalar.linspace(0, 1, 5)'
[]
[strains]
type = Python
expr = 'SR2(torch.tensor([0.01, 0.0, 0.0, 0.0, 0.0, 0.0], dtype=torch.float64).reshape(1, 6) * torch.linspace(0.0, 1.0, 5, dtype=torch.float64).reshape(5, 1))'
[]
[]
[Models]
[elasticity]
type = LinearIsotropicElasticity
coefficients = '200e3 0.3'
coefficient_types = 'YOUNGS_MODULUS POISSONS_RATIO'
[]
[]
[Drivers]
[driver]
type = TransientDriver
model = 'elasticity'
prescribed_time = 'times'
force_SR2_names = 'strain'
force_SR2_values = 'strains'
save_as = 'result.pt'
[]
[]
neml2-run¶
Loads an input file, instantiates the named driver, and calls its
run() method. Use it whenever you want to step a model through a load
history from the shell — from a shell script, a CI job, a parameter
sweep — without writing any Python.
usage: neml2-run [-h] [--device {cpu,cuda}] [--dtype {float64,float32}]
[--load PATH]
input driver
Run a driver from an input file. Trailing positional tokens are treated as HIT
overrides (e.g. 'Models/elasticity/E:=210000').
positional arguments:
input path to the input file
driver name of the driver in the input file
options:
-h, --help show this help message and exit
--device {cpu,cuda} Set torch's default device before loading. Tensors
built by [Tensors] Python expressions inherit this.
Default: cpu.
--dtype {float64,float32}
Set torch's default dtype before loading. NEML2 models
are uniformly float64 by convention; AOTI artifacts
compile-pin their dtype and the runtime rejects silent
coercion, so float64 is the safe default. Override
only if you know the artifact was compiled with
float32.
--load PATH Import a user extension before processing. PATH is
either a filesystem path to a .py file / package
directory, or a dotted module name already importable
from sys.path. Repeatable; modules are imported in the
order given so later modules may depend on names
registered by earlier ones.
Running our shared example: neml2-run itself emits no output on
success and returns a zero exit code, so a clean run looks empty in
CI. (Drivers can still print their own progress.)
$ neml2-run input.i driver
The trailing-token convention lets you tweak scalar options without
editing the file — anything after the two declared positionals is
forwarded to the HIT parser as an override. The same convention works
for neml2-inspect and neml2-compile:
$ neml2-run input.i driver Drivers/driver/save_as:=other.pt
List-valued overrides (e.g. swapping coefficients = '200e3 0.3' for
a different pair) need careful shell quoting around the whole value.
If that gets fiddly, edit the input file directly or use the Python
API.
neml2-inspect¶
Loads a single model from an input file and prints its structural summary: input variables, output variables, parameters, and buffers (each with dtype, device, and shape). Use it any time you want to verify the wiring of a model before driving it — wrong variable names, missing parameters, and unexpected shapes all show up here in one screenful.
usage: neml2-inspect [-h] [--json] [--load PATH] input model
Summarize the structure of a model. Trailing positional tokens are treated as
HIT overrides (e.g. 'Models/elasticity/E:=210000').
positional arguments:
input path to the input file
model name of the model in the input file to inspect
options:
-h, --help show this help message and exit
--json emit a structured JSON description on stdout instead of the
human-readable format
--load PATH Import a user extension before processing. PATH is either a
filesystem path to a .py file / package directory, or a dotted
module name already importable from sys.path. Repeatable;
modules are imported in the order given so later modules may
depend on names registered by earlier ones.
The summary view (default):
$ neml2-inspect input.i elasticity
Model: LinearIsotropicElasticity
Inputs (1):
strain: type=SR2
Outputs (1):
stress: type=SR2
Parameters (2):
E: dtype=torch.float64, device=cpu, shape=[]
nu: dtype=torch.float64, device=cpu, shape=[]
Buffers (0):
The same data is available as a JSON document for programmatic
consumption via --json:
$ neml2-inspect --json input.i elasticity
{"name": "LinearIsotropicElasticity", "inputs": [{"name": "strain", "type": "SR2"}], "outputs": [{"name": "stress", "type": "SR2"}], "parameters": [{"name": "E", "dtype": "torch.float64", "device": "cpu", "shape": []}, {"name": "nu", "dtype": "torch.float64", "device": "cpu", "shape": []}], "buffers": [], "retcode": 0}
...
neml2-syntax¶
Browses the registered-object catalog: every concrete class that can appear in a HIT input file, along with its section, docstring, and full option schema. It is the discovery tool — use it when you don’t yet know which class to instantiate, or when you want to see what options a class accepts without opening the source.
usage: neml2-syntax [-h] [--json JSON] [--section SECTION] [--type TYPE]
[--summary] [--server] [--load PATH]
Extract Python-native object syntax from the native registry as JSON.
options:
-h, --help show this help message and exit
--json JSON redirect JSON output to a file
--section SECTION only emit objects whose input-file section matches
--type TYPE only emit the object whose registered type matches
--summary emit only type, section, source_path, and doc string for
each object
--server run as a JSON server on stdin/stdout
--load PATH Import a user extension before processing. PATH is either
a filesystem path to a .py file / package directory, or a
dotted module name already importable from sys.path.
Repeatable; modules are imported in the order given so
later modules may depend on names registered by earlier
ones.
--summary drops the per-option schema and emits one entry per
registered type — perfect for browsing. Combine with --section to
restrict to a single input-file section:
$ neml2-syntax --section Drivers --summary --json -
[
{
"type": "TransientDriver",
"section": "Drivers",
"doc": "Drive a model over a prescribed time history.\n\nMirrors C++ ``TransientDriver``. See module docstring for differences.",
"source_path": "drivers/TransientDriver.py",
"class_name": ""
},
{
"type": "TransientRegression",
"section": "Drivers",
"doc": "Run a TransientDriver and diff its result against a gold ``.pt`` file.",
"source_path": "drivers/TransientRegression.py",
"class_name": ""
},
{
...
Drop --summary and use --type <Name> for the full schema of one
class — the same data the documentation’s syntax catalog under
Syntax catalog is generated from:
$ neml2-syntax --type LinearIsotropicElasticity --json -
[
{
"type": "LinearIsotropicElasticity",
"section": "Models",
"doc": "Relate elastic strain to stress (or stress to strain, in compliance form)\nfor a linear isotropic material.",
"source_path": "models/solid_mechanics/elasticity/LinearIsotropicElasticity.py",
"class_name": "",
"options": [
{
"name": "strain",
"doc": "Elastic strain",
"ftype": "INPUT",
"required": true,
"type": "SR2",
"value": ""
},
...
neml2-compile¶
Exports a model to an AOT-Inductor package — one or more .pt2
files plus a metadata JSON and a drop-in HIT stub — that can be
loaded later from either Python or C++ without parsing the original
input file. Use it when you want a dependency-light artifact for
inference: shipping a calibrated model to a downstream simulator,
dropping the HIT parser from a deployment image, baking parameters
into a frozen graph for performance.
usage: neml2-compile [-h] (--model NAME | --driver NAME) [--output-dir DIR]
[--device DEVICE [DEVICE ...]]
[--dtype {float64,float32}] [-p NAME] [-d OUT:IN]
[--example-batch-shape SPEC] [--dynamic-batch]
[--no-dynamic-batch] [--load PATH]
INPUT.i
Compile a NEML2 Python-native model to AOTI artifacts (.pt2 + metadata JSON)
and emit a runnable HIT stub.
positional arguments:
INPUT.i HIT input file path.
options:
-h, --help show this help message and exit
--model NAME Compile the named [Models] block. The emitted stub
contains only the AOTIModel shim for this model -- no
[Drivers] block. Pair with a custom driver in a
separate file when you want to run it.
--driver NAME Compile whatever model the named [Drivers] block
targets (via its `model = '...'` field). The stub
keeps the named driver so the result is runnable as-is
via `neml2-run`. Use this when you want a self-
contained, drop-in replacement for the original
driver.
--output-dir DIR Collection dir for the compiled artifacts (default:
./aoti/). Each model lands in <DIR>/<NAME>/<device>/
with a standalone <DIR>/<NAME>_aoti.i stub next to it.
--device DEVICE [DEVICE ...]
Target device(s) for the artifact, baked at export
time. Accepts more than one (e.g. --device cpu cuda):
one complete artifact is emitted per device into a
subfolder named by the device (<output-dir>/cpu/,
<output-dir>/cuda/), ready for a multi-device
dispatcher to load.
--dtype {float64,float32}
Floating-point dtype for the artifact (baked at export
time).
-p NAME, --parameter NAME
Promote a parameter or buffer (fully-qualified dotted
name) to runtime-flexible status -- it becomes a graph
input rather than a baked constant. Repeatable. With
no -p flags the model is fully baked.
-d OUT:IN, --derivative OUT:IN
Compile the derivative (Jacobian / JVP) graph for the
OUT:IN output-input pair. Repeatable. Omit a side to
select all on that side: 'stress:strain' is one pair,
'stress:' all inputs of stress, ':strain' all outputs
wrt strain, ':' all pairs. With no -d flags NO
derivative graphs are compiled (smallest artifact) and
the runtime jvp()/jacobian() raise until recompiled
with -d.
--example-batch-shape SPEC
Example shape used when tracing. Two forms: --example-
batch-shape '(2,)' (uniform across all inputs)
--example-batch-shape strain='(2;100)' --example-
batch-shape temperature='(2,)' (per-variable;
repeatable) The semicolon delimits dynamic-batch from
sub-batch axes -- '(2;100)' means dyn=(2,),
sub=(100,). Overrides [Settings]/example_batch_shape.
--dynamic-batch Compile a dynamic-batch artifact: the leading batch
dim of the example shape becomes a runtime-variable
torch.export Dim. Default unless overridden by
[Settings]/dynamic_batch in the input.
--no-dynamic-batch Pin the batch dim at the example shape. Required when
a baked parameter has rank>=1 and would specialize the
dynamic dim.
--load PATH Import a user extension before processing. PATH is
either a filesystem path to a .py file / package
directory, or a dotted module name already importable
from sys.path. Repeatable; modules are imported in the
order given so later modules may depend on names
registered by earlier ones.
Derivative graphs are opt-in: with no -d flag only forward is
compiled and jvp / jacobian raise at runtime. Request the pairs you
need with -d OUT:IN (e.g. -d stress:strain, or -d : for all).
The end-to-end walkthrough — emitted file layout, loading the
artifact from Python, parameter promotion (-p <name>), and the
trade-offs against eager mode — lives in
Compiled models.
neml2-stub¶
Regenerates .pyi type stubs for every pybind11 extension module in the
installed neml2 package (currently neml2.aoti._aoti). The stubs are
written next to their .so files so that pyright and IDE autocompletion
resolve from ._aoti import Model cleanly.
neml2-stub: generating stubs for neml2.aoti._aoti -> /opt/hostedtoolcache/Python/3.10.20/x64/lib/python3.10/site-packages
usage: pybind11-stubgen [-h] [-o OUTPUT_DIR] [--root-suffix ROOT_SUFFIX]
[--ignore-invalid-expressions REGEX]
[--ignore-invalid-identifiers REGEX]
[--ignore-unresolved-names REGEX]
[--ignore-all-errors]
[--enum-class-locations REGEX:LOC]
[--numpy-array-wrap-with-annotated | --numpy-array-use-type-var | --numpy-array-remove-parameters]
[--print-invalid-expressions-as-is]
[--print-safe-value-reprs REGEX] [--exit-code]
[--dry-run] [--stub-extension EXT]
MODULE_NAME
Generates stubs for specified modules
positional arguments:
MODULE_NAME module name
options:
-h, --help show this help message and exit
-o OUTPUT_DIR, --output-dir OUTPUT_DIR
The root directory for output stubs
--root-suffix ROOT_SUFFIX
Top-level module directory suffix
--ignore-invalid-expressions REGEX
Ignore invalid expressions matching REGEX
--ignore-invalid-identifiers REGEX
Ignore invalid identifiers matching REGEX
--ignore-unresolved-names REGEX
Ignore unresolved names matching REGEX
--ignore-all-errors Ignore all errors during module parsing
--enum-class-locations REGEX:LOC
Locations of enum classes in <enum-class-name-
regex>:<path-to-class> format. Example:
`MyEnum:foo.bar.Baz`
--numpy-array-wrap-with-annotated
Replace numpy/scipy arrays of 'ARRAY_T[TYPE, [*DIMS],
*FLAGS]' format with 'Annotated[ARRAY_T, TYPE,
FixedSize|DynamicSize(*DIMS), *FLAGS]'
--numpy-array-use-type-var
Replace 'numpy.ndarray[numpy.float32[m, 1]]' with
'numpy.ndarray[tuple[M, typing.Literal[1]],
numpy.dtype[numpy.float32]]'
--numpy-array-remove-parameters
Replace 'numpy.ndarray[...]' with 'numpy.ndarray'
--print-invalid-expressions-as-is
Suppress the replacement with '...' of invalid
expressionsfound in annotations
--print-safe-value-reprs REGEX
Override the print-safe check for values matching
REGEX
--exit-code On error exits with 1 and skips stub generation
--dry-run Don't write stubs. Parses module and report errors
--stub-extension EXT The file extension of the generated stubs. Must be
'pyi' (default) or 'py'
This tool is primarily used by CI (between install and pyright) and by the
wheel-building pipeline. Any extra arguments are forwarded verbatim to
pybind11_stubgen.
Where to go next¶
The five scripts are also importable as
neml2.cli.run.main,neml2.cli.inspect.main,neml2.cli.syntax.main,neml2.cli.aoti_compile.main, andneml2.cli.stub.mainif you want to drive them in-process rather than via the shell.The Python-first path —
neml2.load_model(...)+ direct call — is the subject of Running your first model and is usually what you want when you need to inspect or post-process results in memory.Compiled models picks up where
neml2-compileleaves off: loading.pt2artifacts, runtime parameter promotion, and the C++ deployment story.