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.
38