E3J: An Efficient and Open-Source Backend for Euclidean Equivariant Operations on GPU and TPU
Open-source e3j speeds Euclidean-equivariant GPU and TPU kernels, beating cuEquivariance on a MACE simulation.
e3j is an open-source JAX backend for Euclidean-equivariant geometric deep learning on GPU and TPU, using optimized CUDA and Pallas kernels. On a machine-learning interatomic potential workload it reports up to a 34% speed-up over cuEquivariance for water-box NPT simulation with MACE. Tensor-product operations exceed 80% of H100 peak memory bandwidth, and message-passing convolution forward throughput more than doubles earlier backends in many cases. A dedicated Pallas TPU kernel reaches over 80% of TPUv6e memory bandwidth, up to about ten times e3nn-jax, and the library is released under Apache 2.0.
- JAX backend with CUDA and Pallas kernels for Euclidean-equivariant operations.
- Up to 34% faster than cuEquivariance on MACE water-box NPT.
- Tensor products exceed 80% of H100 memory bandwidth.
- TPU kernels exceed 80% of TPUv6e bandwidth, far above e3nn-jax.
- Released on GitHub and PyPI under Apache 2.0.
Full article179 words · extracted from arxiv.org · click to collapse
We present e3j, a fast Euclid-equivariance backend for geometric deep learning applications with JAX bindings for GPU and TPU. Leveraging both optimized CUDA and Pallas kernels and algorithmic improvements, the library achieves state-of-the-art throughput and runtime on both forward and backward paths. On a machine learning interatomic potential (MLIP) use case, it outperforms established backends, measuring up to 34% speed-up over cuEquivariance on water box NPT simulation using MACE, while remaining fully open source. E3j achieves over 80% efficiency over the H100 maximum memory bandwidth on tensor product operations, and in many cases more than doubles throughput of message passing convolutions forward compared to previously available backends. In addition, with the release of dedicated Pallas TPU kernel, e3j opens the possibility of large scale equivariant deep learning workloads on TPU architectures, which has so far been difficult to achieve. Our benchmarks show that e3j also achieves over 80% of a TPUv6e memory bandwidth, up to one order of magnitude more than e3nn-jax. The library is available on GitHub, PyPI and is released under an open source Apache 2.0 license.
Text extracted automatically; images, tables and formatting may be missing. Original: https://arxiv.org/abs/2609.35099