$ show mir-group/allegro
The Allegro package implements the Allegro E(3)-equivariant machine learning interatomic potential as an extension for the NequIP framework. It provides a neural network architecture that predicts energies and forces in atomic systems while being invariant to translations, rotations, and reflections. The model is built using a combination of scalar and tensor feature tracks, interacting through tensor product layers, and is designed to be strictly local with a configurable cutoff radius. The architecture centers around the `AllegroModel` function, which constructs a graph neural network. Input atomic data (positions, types, and optionally cell) is processed through embedding layers to generate initial scalar and tensor features. These features then pass through multiple Allegro layers, each consisting of tensor product operations (via the `Contracter` module) and mixing layers (MLPs). The final features are reduced to per-atom energies, summed to get the total energy, and forces and stresses are computed via automatic differentiation if requested. Key entry points include the `AllegroModel` function and the package's entry point for NequIP integration. The package relies on external dependencies such as NequIP, PyTorch, and e3nn, with optional accelerations via NVIDIA's CuEquivariance library or custom Triton kernels for GPU performance. The codebase demonstrates strong engineering practices, including performance optimizations for the tensor product bottleneck (through both the base `Contracter` implementation and GPU-accelerated variants), a clear modular structure, and comprehensive testing. Reusable components such as the `Contracter` and `MakeWeightedChannels` modules showcase efficient tensor product handling and feature weighting strategies that could be adapted to other equivariant neural network architectures with minimal modification.
The LLM identified 2 files worth highlighting for reuse:
The Contracter class implements an efficient tensor product operation for equivariant neural networks where input irreps have multiplicity 1. It precomputes instructions and optimized contraction paths, handles diagonal cases for performance, and supports both standard and GPU-accelerated variants. This component is self-contained, demonstrates non-obvious skill in optimizing tensor contractions, and can be extracted for use in other projects with minimal modification (requiring only torch and e3nn dependencies).
The MakeWeightedChannels class provides a memory-efficient mechanism to apply weights to tensor features by avoiding explicit construction of large weight matrices. Instead, it uses indexing and broadcasting to compute weighted sums. This self-contained module only depends on torch and e3nn, exhibits careful consideration of computational efficiency, and can be easily reused in other neural network architectures that require similar feature weighting operations.