OpenEquivariance

https://passionlab.github.io/OpenEquivariance/
Admin
OpenEquivariance is a CUDA and HIP kernel generator for the Clebsch-Gordon tensor product, a key...
Tokens:17,184
Snippets:155
Trust Score:9.7
Update:1 month ago
Show doc for...
Tokens:
Raw
### JAX: Install and run JAX tests

Source: https://passionlab.github.io/OpenEquivariance/_sources/tests_and_benchmarks.rst.txt

Installs the JAX backend and runs the JAX-specific tests. Requires PyTorch and e3nn to be installed in the environment for correctness testing.

```bash
git clone https://github.com/PASSIONLab/OpenEquivariance 
cd OpenEquivariance

pip install "./openequivariance[jax]"
pip install "./openequivariance[dev]"
pip install "./openequivariance_extjax" --no-build-isolation

pytest --jax tests/example_test.py
pytest --jax tests/batch_test.py
pytest --jax tests/conv_test.py
pytest --jax tests/vmap_test.py
```

--------------------------------

### Install OpenEquivariance for PyTorch

Source: https://passionlab.github.io/OpenEquivariance/installation

Install the package via pip and verify the installation. The first import triggers a C++ extension build.

```bash
pip install openequivariance
python -c "import openequivariance"
```

--------------------------------

### Install OpenEquivariance for JAX on NVIDIA GPUs

Source: https://passionlab.github.io/OpenEquivariance/installation

Install the JAX-specific package and extension for NVIDIA hardware.

```bash
pip install openequivariance[jax]
pip install openequivariance_extjax --no-build-isolation
```

--------------------------------

### Install Nightly Build for JAX

Source: https://passionlab.github.io/OpenEquivariance/installation

Install the latest development version of OpenEquivariance for JAX.

```bash
pip install "git+https://github.com/PASSIONLab/OpenEquivariance#subdirectory=openequivariance[jax]"
pip install "git+https://github.com/PASSIONLab/OpenEquivariance#subdirectory=openequivariance_extjax --no-build-isolation"
```

--------------------------------

### Benchmark: Install and run all experiments

Source: https://passionlab.github.io/OpenEquivariance/_sources/tests_and_benchmarks.rst.txt

Installs benchmark dependencies and runs the benchmark experiments from the paper, generating plots. Use --limited-memory on GPUs with limited memory and -b to reduce batch size.

```bash
git clone https://github.com/PASSIONLab/OpenEquivariance 
cd OpenEquivariance
pip install -e .[bench] 
python tests/benchmark.py -o outputs/uvu uvu --plot
python tests/benchmark.py -o outputs/uvw uvw --plot
python tests/benchmark.py -o outputs/roofline roofline --plot
python tests/benchmark.py -o outputs/conv conv --plot --data data/molecular_structures
python tests/benchmark.py -o outputs/kahan_conv kahan_conv --data data/molecular_structures/
```

--------------------------------

### PyTorch: Install dev dependencies and run tests

Source: https://passionlab.github.io/OpenEquivariance/_sources/tests_and_benchmarks.rst.txt

Sets up an editable install with development dependencies and runs the test suite for PyTorch. Requires cloning the repository first.

```bash
git clone https://github.com/PASSIONLab/OpenEquivariance 
cd OpenEquivariance
pip install -e "./openequivariance[dev]"
pytest tests/ 
```

--------------------------------

### Install OpenEquivariance for JAX on AMD GPUs

Source: https://passionlab.github.io/OpenEquivariance/installation

Install the JAX-specific package and extension for AMD hardware with RocM support.

```bash
pip install openequivariance[jax]
JAX_HIP=1 pip install openequivariance_extjax --no-build-isolation
```

--------------------------------

### Install Nightly Build for PyTorch

Source: https://passionlab.github.io/OpenEquivariance/installation

Install the latest development version of OpenEquivariance for PyTorch.

```bash
pip install "git+https://github.com/PASSIONLab/OpenEquivariance#subdirectory=openequivariance"
```

--------------------------------

### Install Nightly JAX Version of OpenEquivariance

Source: https://passionlab.github.io/OpenEquivariance/_sources/installation.rst.txt

Installs the nightly JAX version of OpenEquivariance from GitHub. The commented line shows the command for JAX+AMD with JAX_HIP=1.

```bash
pip install "git+https://github.com/PASSIONLab/OpenEquivariance#subdirectory=openequivariance[jax]"
pip install "git+https://github.com/PASSIONLab/OpenEquivariance#subdirectory=openequivariance_extjax --no-build-isolation"

# Use the command below for JAX+AMD
# JAX_HIP=1 pip install "git+https://github.com/PASSIONLab/OpenEquivariance#subdirectory=openequivariance_extjax --no-build-isolation"
```

--------------------------------

### Run correctness tests with PyTorch

Source: https://passionlab.github.io/OpenEquivariance/tests_and_benchmarks

Sets up an editable install with the dev extra and runs the correctness test suite for the PyTorch backend. Requires cloning the repository first.

```shell
git clone https://github.com/PASSIONLab/OpenEquivariance
cd OpenEquivariance
pip install -e "./openequivariance[dev]"
pytest tests/

```

--------------------------------

### Install modified MACE with OpenEquivariance

Source: https://passionlab.github.io/OpenEquivariance/models

Install the modified MACE package that uses OpenEquivariance kernels. First uninstall the standard mace-torch package, then install the modified version from the specified git repository.

```bash
pip uninstall mace-torch
pip install git+https://github.com/vbharadwaj-bk/mace_oeq_integration.git@oeq_experimental

```

--------------------------------

### Benchmark: Example invocations for NVIDIA A5000

Source: https://passionlab.github.io/OpenEquivariance/_sources/tests_and_benchmarks.rst.txt

Example invocations for an NVIDIA A5000 GPU. The roofline slope/peak is incorrect on GPUs other than the NVIDIA A100, and convolution fusion plots require a GPU with at least 40GB of memory.

```bash
python tests/benchmark.py -o outputs/uvu uvu --limited-memory --plot
python tests/benchmark.py -o outputs/uvw uvw -b 25000 --plot
python tests/benchmark.py -o outputs/roofline roofline --plot
python tests/benchmark.py -o outputs/conv conv --data data/molecular_structures --limited-memory
```

--------------------------------

### OLCF Frontier Environment Configuration

Source: https://passionlab.github.io/OpenEquivariance/installation

Environment setup script for OLCF Frontier.

```bash
module load PrgEnv-gnu/8.6.0
module load miniforge3/23.11.0-0
module load rocm/6.4.0
module load craype-accel-amd-gfx90a

for i in $(seq ${CONDA_SHLVL}); do
    conda deactivate
done

conda activate <your-conda-env>
export CC=cc
export CXX=CC
```

--------------------------------

### Replicate paper benchmarks with plots

Source: https://passionlab.github.io/OpenEquivariance/tests_and_benchmarks

Installs the benchmark extra and runs the benchmark experiments from the paper, generating plots for uvu, uvw, roofline, conv, and kahan_conv. The conv and kahan_conv runs require the molecular structures data directory.

```shell
git clone https://github.com/PASSIONLab/OpenEquivariance
cd OpenEquivariance
pip install -e .[bench]
python tests/benchmark.py -o outputs/uvu uvu --plot
python tests/benchmark.py -o outputs/uvw uvw --plot
python tests/benchmark.py -o outputs/roofline roofline --plot
python tests/benchmark.py -o outputs/conv conv --plot --data data/molecular_structures
python tests/benchmark.py -o outputs/kahan_conv kahan_conv --data data/molecular_structures/

```

--------------------------------

### Run correctness tests with JAX

Source: https://passionlab.github.io/OpenEquivariance/tests_and_benchmarks

Installs the JAX extra, the dev extra, and the JAX extension, then runs the JAX-specific correctness tests. Requires PyTorch and e3nn to be present in the environment.

```shell
git clone https://github.com/PASSIONLab/OpenEquivariance
cd OpenEquivariance

pip install "./openequivariance[jax]"
pip install "./openequivariance[dev]"
pip install "./openequivariance_extjax" --no-build-isolation

pytest --jax tests/example_test.py
pytest --jax tests/batch_test.py
pytest --jax tests/conv_test.py
pytest --jax tests/vmap_test.py

```

--------------------------------

### Install modified MACE with oeq

Source: https://passionlab.github.io/OpenEquivariance/_sources/models.rst.txt

Uninstall the standard mace-torch package and install the modified version from the specified GitHub repository branch.

```bash
pip uninstall mace-torch
pip install git+https://github.com/vbharadwaj-bk/mace_oeq_integration.git@oeq_experimental
```

--------------------------------

### Benchmark: Compare against FlashTP

Source: https://passionlab.github.io/OpenEquivariance/_sources/tests_and_benchmarks.rst.txt

Runs the convolution benchmark comparing against FlashTP. Requires FlashTP to be installed via pip. OpenEquivariance shows up to 2x speedup over FlashTP's fused kernels.

```bash
python tests/benchmark.py -o outputs/conv conv --plot --data data/molecular_structures -i cue_unfused oeq_scattersum flashtp cue_fused oeq_det oeq_atomic
```

--------------------------------

### Run benchmarks on limited-memory GPU

Source: https://passionlab.github.io/OpenEquivariance/tests_and_benchmarks

Example invocations for a GPU with limited memory, such as an NVIDIA A5000. Uses the --limited-memory flag to disable expensive tests and -b to reduce batch size; run python tests/benchmark.py --help for the full flag list.

```shell
python tests/benchmark.py -o outputs/uvu uvu --limited-memory --plot
python tests/benchmark.py -o outputs/uvw uvw -b 25000 --plot
python tests/benchmark.py -o outputs/roofline roofline --plot
python tests/benchmark.py -o outputs/conv conv --data data/molecular_structures --limited-memory

```

--------------------------------

### NERSC Perlmutter Environment Configuration

Source: https://passionlab.github.io/OpenEquivariance/installation

Environment setup script for NERSC Perlmutter.

```bash
module load gcc
module load conda

# Deactivate any base environments
for i in $(seq ${CONDA_SHLVL}); do
    conda deactivate
done

conda activate <your-conda-env>
```

--------------------------------

### openequivariance.torch_ext_so_path

Source: https://passionlab.github.io/OpenEquivariance/_sources/api.rst.txt

Function to get the path to the torch extension shared object.

```APIDOC
## openequivariance.torch_ext_so_path

### Description
Returns the path to the torch extension shared object.

```

--------------------------------

### Irreps.slices() example

Source: https://passionlab.github.io/OpenEquivariance/api

Shows the output of calling slices() on an Irreps instance with the string '2x0e + 1e'. The result is a list of slice objects corresponding to the indices for each irrep.

```python
>>> Irreps('2x0e + 1e').slices()
[slice(0, 2, None), slice(2, 5, None)]

```

--------------------------------

### Configure Environment on NERSC Perlmutter

Source: https://passionlab.github.io/OpenEquivariance/_sources/installation.rst.txt

Environment setup for NERSC Perlmutter (NVIDIA A100). Loads GCC and conda modules, deactivates any base environments, and activates your conda environment.

```bash
module load gcc 
module load conda

# Deactivate any base environments
for i in $(seq ${CONDA_SHLVL}); do 
    conda deactivate
done

conda activate <your-conda-env>
```

--------------------------------

### Run full MACE benchmark with multiple backends

Source: https://passionlab.github.io/OpenEquivariance/_sources/models.rst.txt

Execute the benchmark comparing e3nn, cue, and oeq backends, provided sufficient GPU memory is available.

```bash
python tests/mace_driver.py carbon.xyz -o outputs/mace_tests -i e3nn cue oeq
```

--------------------------------

### Benchmark OpenEquivariance with MACE

Source: https://passionlab.github.io/OpenEquivariance/_sources/models.rst.txt

Run the MACE driver using the oeq backend on the carbon.xyz dataset.

```bash
python tests/mace_driver.py carbon.xyz -o outputs/mace_tests -i oeq
```

--------------------------------

### TPProblem constructor

Source: https://passionlab.github.io/OpenEquivariance/api

Creates a TPProblem specification for a CG tensor product. Accepts irreps for inputs and output, a list of instructions, and various configuration parameters for weights, normalization, and layout.

```APIDOC
## TPProblem(irreps_in1, irreps_in2, irreps_out, instructions, in1_var=None, in2_var=None, out_var=None, irrep_normalization='component', path_normalization='element', internal_weights=False, shared_weights=None, label=None, irrep_dtype=numpy.float32, weight_dtype=numpy.float32, layout='mul_ir')

### Description
Specification for a CG tensor product. All parameters from e3nn's o3.TensorProduct are available, along with additional parameters for the types of weights and irreps.

### Parameters
- **irreps_in1** (Irreps) - Required - Irreps for the first CG argument
- **irreps_in2** (Irreps) - Required - Irreps for the second CG argument
- **irreps_out** (Irreps) - Required - Irreps for the output
- **instructions** (List[Any]) - Required - A list of 5-tuples, each of the form (i_in1, i_in2, i_out, has_weight, path_weight). i_in1, i_in2, and i_out each index an Irrep from irreps_in1, irreps_in2, and irreps_in3, respectively. has_weight (True / False) controls whether trainable weights are included for the instruction, and path_weight controls output normalization.
- **irrep_dtype** (type[numpy.generic]) - Optional - Datatype of irrep inputs; one of np.float32 or np.float64. Default: np.float32.
- **weight_dtype** (type[numpy.generic]) - Optional - Datatype of weights; one of np.float32 or np.float64. Default: np.float32.
- **label** (str) - Optional - A name for this problem specification (useful for testing / benchmarking).
- **shared_weights** (bool) - Optional - If True, all elements in a batch of inputs share a common set of weights. If False, each batch element has a unique set of weights. Default: True.
- **internal_weights** (bool) - Optional - Must be False; OpenEquivariance does not support internal weights. Default: False.
- **irrep_normalization** (str) - Optional - One of ["component", "norm", "none"]. Default: "component".
- **path_normalization** (str) - Optional - One of ["element", "path", "none"]. Default: "element".
- **layout** (str) - Optional - One of ["mul_ir", "ir_mul"], giving the layout of irreps for all inputs and outputs. Default: "mul_ir".
- **in1_var** (List[float] | None) - Optional - Description not provided in source.
- **in2_var** (List[float] | None) - Optional - Description not provided in source.
- **out_var** (List[float] | None) - Optional - Description not provided in source.
```

--------------------------------

### Configure JAX without PyTorch

Source: https://passionlab.github.io/OpenEquivariance/installation

Set the OEQ_NOTORCH environment variable to avoid importing PyTorch when using JAX.

```bash
export OEQ_NOTORCH=1
python -c "import openequivariance.jax"
```

--------------------------------

### openequivariance.TPProblem

Source: https://passionlab.github.io/OpenEquivariance/_sources/api.rst.txt

Configuration object required by TensorProduct and TensorProductConv. Constructor is almost identical to o3.TensorProduct.

```APIDOC
## openequivariance.TPProblem

### Description
Configuration object required by TensorProduct and TensorProductConv. Constructor is almost identical to o3.TensorProduct.

```

--------------------------------

### TPProblem.clone()

Source: https://passionlab.github.io/OpenEquivariance/api

Creates a clone of the TPProblem instance. The method signature is documented but no additional details are provided in the source.

```APIDOC
## clone()

### Description
Creates a clone of the TPProblem instance. The method signature is documented but no additional details are provided in the source.
```

--------------------------------

### Import SymmetricContraction from openequivariance

Source: https://passionlab.github.io/OpenEquivariance/_sources/supported_ops.rst.txt

Import the beta symmetric contraction implementation, which requires e3nn as a dependency and is not exposed at the package toplevel.

```python
from openequivariance._torch.symmetric_contraction import SymmetricContraction as OEQSymmetricContraction
```

--------------------------------

### openequivariance.jax.TensorProductConv

Source: https://passionlab.github.io/OpenEquivariance/_sources/api.rst.txt

JAX version of TensorProductConv, behaves identically to PyTorch counterpart. Provides methods for forward computation and weight reordering.

```APIDOC
## openequivariance.jax.TensorProductConv

### Description
JAX version of TensorProductConv, behaves identically to PyTorch counterpart.

### Methods
- **forward** - Performs the fused tensor product and convolution operation.
- **reorder_weights_from_e3nn** - Reorders weights from e3nn format.
- **reorder_weights_to_e3nn** - Reorders weights to e3nn format.

```

--------------------------------

### Build Extension with Alternate Compiler

Source: https://passionlab.github.io/OpenEquivariance/installation

Set the CC and CXX environment variables to use a specific compiler for the C++ extension build.

```bash
export CC=/path/to/your/gcc
export CXX=/path/to/your/g++
python -c "import openequivariance"
```

--------------------------------

### openequivariance.jax.TensorProduct

Source: https://passionlab.github.io/OpenEquivariance/_sources/api.rst.txt

JAX version of TensorProduct, behaves identically to PyTorch counterpart. Provides methods for forward computation and weight reordering.

```APIDOC
## openequivariance.jax.TensorProduct

### Description
JAX version of TensorProduct, behaves identically to PyTorch counterpart.

### Methods
- **forward** - Performs the tensor product operation.
- **reorder_weights_from_e3nn** - Reorders weights from e3nn format.
- **reorder_weights_to_e3nn** - Reorders weights to e3nn format.

```

--------------------------------

### TensorProduct.to(*args, **kwargs)

Source: https://passionlab.github.io/OpenEquivariance/api

Moves the module and its parameters to the specified device and/or dtype. See torch.nn.Module.to().

```APIDOC
## TensorProduct.to(*args, **kwargs)

### Description
Moves the module and its parameters to the specified device and/or dtype. See torch.nn.Module.to().

### Method
to

### Parameters
#### Parameters
- **args** - Required - Positional arguments passed to torch.nn.Module.to().
- **kwargs** - Required - Keyword arguments passed to torch.nn.Module.to().
```

--------------------------------

### openequivariance.TensorProduct

Source: https://passionlab.github.io/OpenEquivariance/api

Drop-in replacement for o3.TensorProduct from e3nn. Supports forward, backward, and double-backward passes using JIT-compiled kernels. Initialization fails if there are no visible GPUs or the tensor product specification is unsupported.

```APIDOC
## openequivariance.TensorProduct

### Description
Drop-in replacement for `o3.TensorProduct` from e3nn. Supports forward, backward, and double-backward passes using JIT-compiled kernels. Initialization fails if there are no visible GPUs or the tensor product specification is unsupported.

### Constructor Parameters
- **problem** (TPProblem) - Required - Specification of the tensor product.
- **use_opaque** - Deprecated.

### Methods

#### forward(x, y, W)
Computes W(x⊗CGy), identical to `o3.TensorProduct.forward`.

**Parameters:**
- **x** (torch.Tensor) - Required - Tensor of shape `[batch_size, problem.irreps_in1.dim()]`, datatype `problem.irrep_dtype`.
- **y** (torch.Tensor) - Required - Tensor of shape `[batch_size, problem.irreps_in2.dim()]`, datatype `problem.irrep_dtype`.
- **W** (torch.Tensor) - Required - Tensor of datatype `problem.weight_dtype` and shape:
  - `[batch_size, problem.weight_numel]` if `problem.shared_weights=False`
  - `[problem.weight_numel]` if `problem.shared_weights=True`

**Returns:**
- Tensor of shape `[batch_size, problem.irreps_out.dim()]`, datatype `problem.irrep_dtype`.

**Return type:** torch.Tensor

#### reorder_weights_from_e3nn(weights, has_batch_dim=True)
Reorders weights from `e3nn` canonical order to the order used by `oeq`.

**Parameters:**
- **weights** - Required - Weights in `e3nn` canonical order, either an np.ndarray, torch.Tensor or JAX array. Tensor of dimensions `[B, problem.weight_numel]` when `has_batch_dim=True`, otherwise of dimensions `[problem.weight_numel]`.
- **has_batch_dim** (bool) - Optional - If `True`, treats the first dimension of weights as a batch dimension. Default: `True`.

**Returns:**
- Weights in `oeq` order. Output type is identical to input.

#### reorder_weights_to_e3nn(weights, has_batch_dim=True)
Reorders weights from `oeq` canonical order to the order used by `e3nn`.

**Parameters:**
- **weights** - Required - Weights in `oeq` canonical order, either a np.ndarray, torch.Tensor or JAX array. Tensor of dimensions `[B, problem.weight_numel]` when `has_batch_dim=True`, otherwise of dimensions `[problem.weight_numel]`.
- **has_batch_dim** (bool) - Optional - If `True`, treats the first dimension of weights as a batch dimension. Default: `True`.

**Returns:**
- Weights in `e3nn` order. Output type is identical to input.

#### to(*args, **kwargs)
See torch.nn.Module.to().
```

--------------------------------

### TPProblem.weight_range_and_shape_for_instruction(instruction)

Source: https://passionlab.github.io/OpenEquivariance/api

Returns the weight range and shape for a given instruction index. The method takes an integer instruction index and returns a tuple of (int, int, tuple).

```APIDOC
## weight_range_and_shape_for_instruction(instruction)

### Description
Returns the weight range and shape for a given instruction index.

### Parameters
- **instruction** (int) - Required - The instruction index.

### Return Type
Tuple[int, int, tuple]
```

--------------------------------

### Load a saved module with torch.load

Source: https://passionlab.github.io/OpenEquivariance/_sources/supported_ops.rst.txt

Call import openequivariance before torch.load to restore a saved TensorProduct or TensorProductConv module.

```python
import torch
import openequivariance
module = torch.load("my_module_with_tp.pt")
```

--------------------------------

### openequivariance.TensorProduct

Source: https://passionlab.github.io/OpenEquivariance/_sources/api.rst.txt

Class that replaces o3.TensorProduct from e3nn. Initialization triggers JIT compilation of a custom kernel. Provides methods for forward computation, weight reordering, and conversion.

```APIDOC
## openequivariance.TensorProduct

### Description
Replaces o3.TensorProduct from e3nn. Initialization triggers JIT compilation of a custom kernel.

### Methods
- **forward** - Performs the tensor product operation.
- **reorder_weights_from_e3nn** - Reorders weights from e3nn format.
- **reorder_weights_to_e3nn** - Reorders weights to e3nn format.
- **to** - Moves the module to a specified device/dtype.

```

--------------------------------

### Check C++ Compiler Version

Source: https://passionlab.github.io/OpenEquivariance/installation

Verify the compiler version if the C++ extension fails to compile.

```bash
c++ --version
```

--------------------------------

### TensorProduct.forward(X, Y, W, rows, cols, sender_perm=None)

Source: https://passionlab.github.io/OpenEquivariance/api

Computes the tensor product of input features X and Y with weights W, using the adjacency matrix defined by rows and cols. Returns the output feature tensor.

```APIDOC
## TensorProduct.forward(X, Y, W, rows, cols, sender_perm=None)

### Description
Computes the tensor product of input features X and Y with weights W, using the adjacency matrix defined by rows and cols. Returns the output feature tensor.

### Method
forward

### Parameters
#### Parameters
- **X** (torch.Tensor) - Required - Tensor of shape [|V|, problem.irreps_in1.dim()], datatype problem.irrep_dtype.
- **Y** (torch.Tensor) - Required - Tensor of shape [|E|, problem.irreps_in1.dim()], datatype problem.irrep_dtype.
- **W** (torch.Tensor) - Required - Tensor of datatype problem.weight_dtype and shape [|E|, problem.weight_numel] if problem.shared_weights=False, or [problem.weight_numel] if problem.shared_weights=True.
- **rows** (torch.Tensor) - Required - Tensor of shape [|E|] with row indices for each nonzero in the adjacency matrix, datatype torch.int64. Must be row-major sorted along with cols when deterministic=True.
- **cols** (torch.Tensor) - Required - Tensor of shape [|E|] with column indices for each nonzero in the adjacency matrix, datatype torch.int64.
- **sender_perm** (torch.Tensor | None) - Optional - Tensor of shape [|E|] and torch.int64 datatype containing a permutation that transposes the adjacency matrix nonzeros from row-major to column-major order. Must be provided when deterministic=True.

### Returns
Tensor of shape [|V|, problem.irreps_out.dim()], datatype problem.irrep_dtype.

### Return Type
torch.Tensor
```

--------------------------------

### TensorProduct.forward(X, Y, W)

Source: https://passionlab.github.io/OpenEquivariance/api

Performs the forward pass of the tensor product operation in JAX. Takes input tensors X, Y, and weights W, and returns the output tensor.

```APIDOC
## TensorProduct.forward(X, Y, W)

### Description
Performs the forward pass of the tensor product operation in JAX. Takes input tensors X, Y, and weights W, and returns the output tensor.

### Method
forward

### Parameters
- **X** (jax.numpy.ndarray) - Required - Input tensor X.
- **Y** (jax.numpy.ndarray) - Required - Input tensor Y.
- **W** (jax.numpy.ndarray) - Required - Weights tensor.

### Return Type
jax.numpy.ndarray
```

--------------------------------

### Run export tests with pytest

Source: https://passionlab.github.io/OpenEquivariance/_sources/supported_ops.rst.txt

Run the export tests to demo C++ model exports. Note: AOTInductor test may fail without a PyTorch Nightly build past 4/10/2025.

```bash
pytest tests/export_test.py
```

--------------------------------

### Irreps.simplify()

Source: https://passionlab.github.io/OpenEquivariance/api

Simplifies the Irreps representation.

```APIDOC
## Irreps.simplify()

### Description
Simplifies the Irreps representation.

### Returns
- **Irreps** - Simplified Irreps.
```

--------------------------------

### Irreps.ls

Source: https://passionlab.github.io/OpenEquivariance/api

Property that returns the list of l values for each irrep.

```APIDOC
## Irreps.ls

### Description
Property that returns the list of l values for each irrep.

### Returns
- **List[int]** - List of l values.
```

--------------------------------

### openequivariance.TensorProductConv

Source: https://passionlab.github.io/OpenEquivariance/_sources/api.rst.txt

Class that fuses the CG tensor product with a subsequent graph convolution. Initialization triggers JIT compilation of a custom kernel. Provides methods for forward computation, weight reordering, and conversion.

```APIDOC
## openequivariance.TensorProductConv

### Description
Fuses the CG tensor product with a subsequent graph convolution. Initialization triggers JIT compilation of a custom kernel.

### Methods
- **forward** - Performs the fused tensor product and convolution operation.
- **reorder_weights_from_e3nn** - Reorders weights from e3nn format.
- **reorder_weights_to_e3nn** - Reorders weights to e3nn format.
- **to** - Moves the module to a specified device/dtype.

```

--------------------------------

### Irreps.sort()

Source: https://passionlab.github.io/OpenEquivariance/api

Sorts the irreps.

```APIDOC
## Irreps.sort()

### Description
Sorts the irreps.

### Returns
- **Irreps** - Sorted Irreps.
```

--------------------------------

### Irreps.spherical_harmonics(lmax, p=-1)

Source: https://passionlab.github.io/OpenEquivariance/api

Static method that returns an Irreps object for spherical harmonics up to lmax.

```APIDOC
## Irreps.spherical_harmonics(lmax, p=-1)

### Description
Static method that returns an Irreps object for spherical harmonics up to lmax.

### Parameters
- **lmax** (int) - Required - Maximum l value.
- **p** (int) - Optional - Parity, default -1.

### Returns
- **Irreps** - Irreps for spherical harmonics.
```

--------------------------------

### openequivariance.Irreps

Source: https://passionlab.github.io/OpenEquivariance/_sources/api.rst.txt

Irreps class identical to e3nn.o3.Irreps. Can be freely mixed with e3nn objects.

```APIDOC
## openequivariance.Irreps

### Description
Irreps class identical to e3nn.o3.Irreps. Can be freely mixed with e3nn objects.

```

--------------------------------

### TensorProductConv.forward(X, Y, W, rows, cols, sender_perm=None)

Source: https://passionlab.github.io/OpenEquivariance/api

Performs the forward pass of the tensor product convolution operation in JAX. Takes input tensors X, Y, weights W, and index arrays rows and cols, with optional sender_perm.

```APIDOC
## TensorProductConv.forward(X, Y, W, rows, cols, sender_perm=None)

### Description
Performs the forward pass of the tensor product convolution operation in JAX. Takes input tensors X, Y, weights W, and index arrays rows and cols, with optional sender_perm.

### Method
forward

### Parameters
- **X** (jax.numpy.ndarray) - Required - Input tensor X.
- **Y** (jax.numpy.ndarray) - Required - Input tensor Y.
- **W** (jax.numpy.ndarray) - Required - Weights tensor.
- **rows** (jax.numpy.ndarray) - Required - Row indices for the convolution.
- **cols** (jax.numpy.ndarray) - Required - Column indices for the convolution.
- **sender_perm** (jax.numpy.ndarray | None) - Optional - Permutation for sender indices. Default: None.

### Return Type
jax.numpy.ndarray
```

--------------------------------

### TensorProduct.reorder_weights_to_e3nn(weights, has_batch_dim=True)

Source: https://passionlab.github.io/OpenEquivariance/api

Reorders weights from the OpenEquivariance format to the e3nn format. See oeq.TensorProduct.reorder_weights_to_e3nn().

```APIDOC
## TensorProduct.reorder_weights_to_e3nn(weights, has_batch_dim=True)

### Description
Reorders weights from the OpenEquivariance format to the e3nn format. See oeq.TensorProduct.reorder_weights_to_e3nn().

### Method
reorder_weights_to_e3nn

### Parameters
#### Parameters
- **weights** - Required - Weights to reorder.
- **has_batch_dim** (bool) - Optional - Whether the weights have a batch dimension. Default is True.
```

--------------------------------

### openequivariance.torch_to_oeq_dtype

Source: https://passionlab.github.io/OpenEquivariance/_sources/api.rst.txt

Function to convert torch dtype to OpenEquivariance dtype.

```APIDOC
## openequivariance.torch_to_oeq_dtype

### Description
Converts a torch dtype to an OpenEquivariance dtype.

```

--------------------------------

### openequivariance.jax.transpose_irreps

Source: https://passionlab.github.io/OpenEquivariance/_sources/api.rst.txt

JAX version of transpose_irreps function.

```APIDOC
## openequivariance.jax.transpose_irreps

### Description
Transposes irreps (JAX version).

```

--------------------------------

### openequivariance.transpose_irreps

Source: https://passionlab.github.io/OpenEquivariance/_sources/api.rst.txt

Function to transpose irreps.

```APIDOC
## openequivariance.transpose_irreps

### Description
Transposes irreps.

```

--------------------------------

### transpose_irreps(array, irreps, src_layout, dst_layout)

Source: https://passionlab.github.io/OpenEquivariance/api

Transposes irrep-packed feature arrays between mul_ir and ir_mul layouts. Uses differentiable JAX operations, so gradients propagate through the transpose.

```APIDOC
## transpose_irreps(array, irreps, src_layout, dst_layout)

### Description
Transposes irrep-packed feature arrays between mul_ir and ir_mul layouts. Uses differentiable JAX operations, so gradients propagate through the transpose.

### Method
transpose_irreps

### Parameters
- **array** (jax.Array) - Required - Input feature array with shape [..., irreps.dim].
- **irreps** (Irreps) - Required - Irreps specification describing how the trailing feature dimension is partitioned into irrep blocks.
- **src_layout** (str) - Required - Source layout. Must be either "mul_ir" or "ir_mul".
- **dst_layout** (str) - Required - Destination layout. Must be either "mul_ir" or "ir_mul".

### Returns
Array in dst_layout with the same shape, dtype, and device as array. If src_layout == dst_layout, returns a copy of array.

### Return Type
jax.Array

### Raises
- **ValueError** - If src_layout or dst_layout is not one of "mul_ir" or "ir_mul".
```

--------------------------------

### TensorProduct.reorder_weights_from_e3nn(weights, has_batch_dim=True)

Source: https://passionlab.github.io/OpenEquivariance/api

Reorders weights from the e3nn format to the OpenEquivariance format. See oeq.TensorProduct.reorder_weights_from_e3nn().

```APIDOC
## TensorProduct.reorder_weights_from_e3nn(weights, has_batch_dim=True)

### Description
Reorders weights from the e3nn format to the OpenEquivariance format. See oeq.TensorProduct.reorder_weights_from_e3nn().

### Method
reorder_weights_from_e3nn

### Parameters
#### Parameters
- **weights** - Required - Weights to reorder.
- **has_batch_dim** (bool) - Optional - Whether the weights have a batch dimension. Default is True.
```

--------------------------------

### Irreps.slices()

Source: https://passionlab.github.io/OpenEquivariance/api

Returns a list of slices corresponding to indices for each irrep.

```APIDOC
## Irreps.slices()

### Description
Returns a list of slices corresponding to indices for each irrep.

### Returns
- **List[slice]** - List of slices.

### Example
```
>>> Irreps('2x0e + 1e').slices()
[slice(0, 2, None), slice(2, 5, None)]
```
```

--------------------------------

### Irreps.regroup()

Source: https://passionlab.github.io/OpenEquivariance/api

Regroups the irreps into a simplified form.

```APIDOC
## Irreps.regroup()

### Description
Regroups the irreps into a simplified form.

### Returns
- **Irreps** - Regrouped Irreps.
```

--------------------------------

### TensorProductConv.reorder_weights_from_e3nn(weights, has_batch_dim=True)

Source: https://passionlab.github.io/OpenEquivariance/api

Reorders weights from e3nn canonical order to the order used by oeq. See oeq.TensorProduct.reorder_weights_from_e3nn().

```APIDOC
## TensorProductConv.reorder_weights_from_e3nn(weights, has_batch_dim=True)

### Description
Reorders weights from e3nn canonical order to the order used by oeq. See oeq.TensorProduct.reorder_weights_from_e3nn().

### Method
reorder_weights_from_e3nn

### Parameters
- **weights** - Required - Weights in e3nn canonical order, either an np.ndarray, torch.Tensor or JAX array. Tensor of dimensions [B, problem.weight_numel] when has_batch_dim=True, otherwise of dimensions [problem.weight_numel].
- **has_batch_dim** (bool) - Optional - If True, treats the first dimension of weights as a batch dimension. Default: True.

### Returns
Weights in oeq order. Output type is identical to input.
```
Tokens:
Raw