E3J: An Efficient and Open-Source Backend for Euclidean Equivariant Operations on GPU and TPU
E3j, a fast Euclid-equivariance backend for geometric deep learning applications with JAX bindings for GPU and TPU architectures, is presented, which opens the possibility of large scale equivariant deep learning workloads on TPU architectures, which has so far been difficult to achieve.
Olivier Peltre, A. Picard, Alexandre Frederic Julien Pichard et al.
· 0 citations