cuequivariance-ops-jax

Introduction

cuequivariance_ops_jax provides CUDA kernels for the cuEquivariance project's JAX components. As such, it contains JAX bindings to optimized kernels that cuEquivariance's operations map down to. In general, we advice that you access those kernels through cuEquivariance, but you may also find them useful on their own.

Installation

Please install using pip install cuequivariance-ops-jax-cu12 or pip install cuequivariance-ops-jax-cu13 (depending on the CUDA toolkit you wish to use).

Documentation

For detailed usage information of the kernels, please refer to the doc-strings in their respective modules. For higher-level documentations, refer to cuEquivariance.

Usage

You can import the library from python:

import cuequivariance_ops_jax

Support and Feedback

Please contact the cuEquivariance developers for any issues you might encounter.