Track synthesis, materials and catalysis as it is published.
Search papers, labs, and topics across Lattice.
The authors developed e3j, an open-source JAX library targeting Euclidean equivariant operations through algorithmic improvements and specialized CUDA and Pallas kernels for both GPUs and TPUs. Geometric deep learning workloads like molecular simulations have historically struggled with memory bandwidth bottlenecks and a near-total absence of optimized TPU runtimes for tensor products. Across benchmarks, e3j saturates over 80% of peak memory bandwidth on both H100 GPUs and TPUv6e, outperforming cuEquivariance by up to 34% on MACE simulations and beating e3nn-jax by up to an order of magnitude on TPUs.
Equivariant tensor products can now saturate over 80% of peak memory bandwidth across both H100 GPUs and TPUv6e, beating cuEquivariance by up to 34% within an entirely open-source JAX framework.
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.