e3nn: Euclidean Neural Networks
It provides a generalized tool for researchers and practitioners working on 3D machine learning tasks, such as molecular modeling or computer vision, but is incremental as it unifies and extends prior methods.
The paper introduces e3nn, a framework for building E(3) equivariant neural networks that handle 3D geometric data with predictable transformations under coordinate changes, enabling efficient implementation of various existing equivariant network architectures.
We present e3nn, a generalized framework for creating E(3) equivariant trainable functions, also known as Euclidean neural networks. e3nn naturally operates on geometry and geometric tensors that describe systems in 3D and transform predictably under a change of coordinate system. The core of e3nn are equivariant operations such as the TensorProduct class or the spherical harmonics functions that can be composed to create more complex modules such as convolutions and attention mechanisms. These core operations of e3nn can be used to efficiently articulate Tensor Field Networks, 3D Steerable CNNs, Clebsch-Gordan Networks, SE(3) Transformers and other E(3) equivariant networks.