Dependencies

Runtime

A PyPI wheel install pulls in three Python packages:

  • pyzag — adapter that lets NEML2 models participate in PyTorch training loops.

  • nmhit — HIT input-file parser.

Pip handles all three automatically. Every other library NEML2 needs at runtime is bundled inside the wheel or resolved from the installed torch package automatically.

ABI compatibility between the wheel and your installed torch is verified for the version range listed in Torch compatibility.

Source-build dependencies

Building NEML2 from source — for contributors, or for build flavors the wheels don’t ship; see Building from source — adds:

  • A C++17 compiler.

  • Python >= 3.10

  • PyTorch importable in the same Python environment used to drive the build (CMake discovers libtorch via the installed torch package by default).

  • Optionally, MPI when configuring with -DNEML2_MPI=ON to enable distributed dispatching.

Everything else (nlohmann_json, etc.) is vendored as a git submodule under contrib/ and pulled in automatically.

GPU acceleration

Accelerator support comes from torch, not from NEML2 directly. The torch that pip install neml2 pulls in by default is whatever your platform’s default PyPI index ships — currently the CUDA-enabled build on x86_64 Linux. If you’re on a platform whose default wheel is CPU-only, or you need a different variant (specific CUDA runtime, ROCm, nightly, …), install that torch first using the selector at pytorch.org/get-started, then run pip install neml2. NEML2 transparently runs models on whichever device the input tensors live on.

CUDA AOTI export

neml2-compile targeting CUDA needs nvcc on top of torch’s bundled CUDA runtime. pip install nvidia-cuda-nvcc ships it as a regular Python package; point CUDA_HOME at its install root and add bin/ to PATH. A system-wide CUDA toolkit (apt install nvidia-cuda-toolkit, conda’s cudatoolkit-dev, …) works too if you already have one. CPU-only AOTI compile needs none of this. And if you do attempt a CUDA compile without nvcc on PATH, neml2-compile stops with a clear error message that includes the install recipe above.