E3J: An Efficient and Open-Source Backend for Euclidean Equivariant Operations on GPU and TPU

Olivier Peltre, Armand Picard, Adrien Pichard +6

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…