Installation

Requirements

TACE requires Python 3.9 or newer and PyTorch 2.4 or newer. AOTInductor export additionally requires torch>=2.13. We recommend installing TACE in a clean environment:

micromamba create -n tace python=3.13 -y
micromamba activate tace

Install TACE

Install the latest release from PyPI:

pip install tace

To install the current source tree instead:

pip install git+https://github.com/xvzemin/tace.git@main

The core installation uses the standard e3nn implementation. Acceleration libraries and simulation interfaces are optional and can be installed independently as described below. When working from a source checkout, replace tace[extra] with .[extra] in the commands.

Core CPU inference does not require CUDA or a C++ compiler. On Windows and macOS, start with device="cpu" and leave the CUDA acceleration backends disabled. EQX CUDA kernels require an NVIDIA GPU, the CUDA toolkit and a compatible C++ compiler (MSVC on Windows). They are not available on macOS. The LAMMPS interface additionally requires a CUDA Kokkos ML-IAP build; the host-only LAMMPS backend is not supported. See LAMMPS ML-IAP.

OpenEquivariance (OEQ)

OEQ provides optimized CUDA or HIP equivariant kernels:

pip install "tace[oeq]"

Important

When OEQ is used together with AOTInductor export or deployment, TACE requires openequivariance>=0.6.4. Upgrade OEQ before exporting the AOTI package:

# OEQ used together with AOTI
pip install "openequivariance>=0.6.4"

Enable it before constructing or loading a configurable model:

export TACE_USE_OEQ=1

cuEquivariance (CUEQ)

Install the package matching the CUDA major version used by PyTorch. CUDA 12 and CUDA 13 use different kernel packages:

# CUDA 12
pip install "tace[cueq12]"

# CUDA 13
pip install "tace[cueq13]"

Check torch.version.cuda if the correct CUDA variant is unclear, then enable the backend with:

python -c "import torch; print(torch.version.cuda)"
export TACE_USE_CUE=1

EquiTorch (EQT)

The EQT implementation used by TACE is bundled with TACE, so ordinary EQT usage does not require installing a separate EquiTorch package:

export TACE_USE_EQT=1

The sparse higher-order product path uses torch-scatter when available and otherwise falls back to native PyTorch reductions. To install the optional extension, select a wheel matching the PyTorch and CUDA versions in the environment. For example, for PyTorch 2.11 and CUDA 13.0:

pip install torch-scatter \
  -f https://data.pyg.org/whl/torch-2.11.0+cu130.html

Use the PyTorch Geometric installation guide to select a different PyTorch or CUDA wheel.

Latent Ewald Summation (LES)

LES is an optional external dependency used by TACE-LES for long-range interactions. Install TACE first, then install the upstream LES (v0.2.0).

pip install git+https://github.com/ChengUCB/les.git@v0.2.0

Verify that TACE can import the backend:

python -c "from les import Les; print('LES is available')"

See the TACE-LES tutorial for model configuration, supported latent sources, outputs, and current compatibility limitations.

TorchSim

Install the optional TorchSim interface with:

pip install "tace[torchsim]"

Important

TACE requires torch-sim-atomistic>=0.6.1. The recommended version is 0.6.2:

pip install "torch-sim-atomistic==0.6.2"

See the TorchSim Calculator tutorial for calculator usage.

NValCHEMI

The NValCHEMI interface is optional and requires Python >= 3.11. Install TACE together with the interface dependencies using:

pip install "tace[nvalchemi]"

For a source checkout, use:

pip install ".[nvalchemi]"

The extra installs nvalchemi-toolkit and its required nvalchemi-toolkit-ops dependency.

The two upstream NVIDIA repositories are provided for reference:

• nvalchemi-toolkit

• nvalchemi-toolkit-ops

Time-reversal e3nn (T-e3nn)

T-e3nn is a time-reversal extension of e3nn, developed by Hongyu Yu et al.

Project repository: T-e3nn

Since the original T-e3nn was developed based on an older version of e3nn, we have migrated and adapted its implementation to the latest version of e3nn for compatibility with the latest APIs and features.

pip install --force-reinstall --no-deps \
  "e3nn @ git+https://github.com/xvzemin/e3nn.git@time-reversal"

The package name and Python import remain e3nn.

TACE detects this capability at runtime; no model option is required. Pure O(3) models then attach time-reversal parity to magnetic moments, magnetic fields, and their tensor-product paths automatically.

Time-reversal models use the time-reversal e3nn implementation for global representation metadata and coupling rules. Native EQX O(2) operators and compatible fused kernels preserve these labels. EQT, CUEQ, and OEQ operators that do not support time-odd irreps reject them when selected; support is checked per operator rather than by removing time-odd coupling paths.

EquivariantX

EquivariantX is currently bundled with TACE and does not require a separate installation for TACE users. Independent installation currently supports source builds only and does not require installing TACE. A standalone package release is planned once the library is fully mature. To install from source:

git clone https://github.com/xvzemin/tace.git
pip install ./tace/eqx

The library imports as eqx and depends on PyTorch >= 2.4 and e3nn >= 0.4.4, but not on TACE or PyG. Its scope covers native O(2) operators, e3nn-compatible O(3)/O(2) frame conversion, and fused CUDA convolutions. Install './tace/eqx[cuda]' and provide a CUDA toolkit to use the fused backend. See Installation and Fused convolutions for installation and operator parameters. When using e3nn 0.4.x with recent PyTorch, import eqx before e3nn.o3 so its packaged constants are loaded in the scoped compatibility context.

Acceleration Selection

Enabled backends are selected per supported operator in the order EQX, OEQ, EQT, CUEQ. Multiple backends may be enabled together, for example OEQ convolutions with EQT product-basis operations. AOTI is a separate compilation and deployment layer. See the Acceleration for backend selection, Python interfaces, compilation, and AOTI export.