LunaEquivariance 0.1.0__tar.gz

lunaequivariance-0.1.0/LICENSE

@@ -0,0 +1,29 @@
+BSD 3-Clause License
+
+Copyright (c) 2026, HPC-AI-Team
+All rights reserved.
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met:
+
+1. Redistributions of source code must retain the above copyright notice, this
+   list of conditions and the following disclaimer.
+
+2. Redistributions in binary form must reproduce the above copyright notice,
+   this list of conditions and the following disclaimer in the documentation
+   and/or other materials provided with the distribution.
+
+3. Neither the name of the copyright holder nor the names of its contributors
+   may be used to endorse or promote products derived from this software
+   without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

lunaequivariance-0.1.0/LunaEquivariance.egg-info/PKG-INFO

@@ -0,0 +1,165 @@
+Metadata-Version: 2.4
+Name: LunaEquivariance
+Version: 0.1.0
+Summary: High-performance SO(3)-equivariant operators for graph neural networks.
+Author: HPC-AI-Team
+License-Expression: BSD-3-Clause
+Project-URL: Homepage, https://github.com/HPC-AI-Team/LunaEquivariance
+Project-URL: Repository, https://github.com/HPC-AI-Team/LunaEquivariance
+Project-URL: Documentation, https://hpc-ai-team.github.io/LunaEquivariance/
+Project-URL: Issues, https://github.com/HPC-AI-Team/LunaEquivariance/issues
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: C++
+Classifier: Operating System :: POSIX :: Linux
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+License-File: src/third_party/cutlass/LICENSE.txt
+Requires-Dist: torch<3,>=2.8
+Requires-Dist: numpy
+Requires-Dist: ninja
+Requires-Dist: e3nn<0.7,>=0.6
+Provides-Extra: test
+Requires-Dist: pytest<9,>=7; extra == "test"
+Provides-Extra: dev
+Requires-Dist: LunaEquivariance[test]; extra == "dev"
+Requires-Dist: mypy<2,>=1.10; extra == "dev"
+Requires-Dist: build<2,>=1; extra == "dev"
+Dynamic: license-file
+
+# LunaEquivariance
+
+<p align="center">
+  <img src="docs/_static/logo.svg" alt="LunaEquivariance" width="480">
+</p>
+
+High-performance SO(3)-equivariant operators for graph neural networks.
+
+## Features
+
+- **Performance:** order-of-magnitude speedups across a range of operators.
+- **Operator coverage:** equivariant layers, fused graph convolutions and attention.
+- **Differentiation and precision:** derivatives up to third order, with FP32/FP64 support.
+- **PyTorch integration:** easy integration into PyTorch models, with `torch.compile` support.
+
+## Installation and user guides
+
+Install PyTorch first, then install LunaEquivariance:
+
+```bash
+pip install LunaEquivariance
+```
+
+See the [documentation](https://hpc-ai-team.github.io/LunaEquivariance/) for installation requirements and user guides.
+
+## Supported operators
+
+| Category | Operators |
+|---|---|
+| Tensor products | [`TensorProduct`](https://hpc-ai-team.github.io/LunaEquivariance/operators/tensor_product.html), [`TensorProductConv`](https://hpc-ai-team.github.io/LunaEquivariance/operators/tensor_product_conv.html) |
+| Geometry and linear layers | [`o3.SphericalHarmonics`](https://hpc-ai-team.github.io/LunaEquivariance/operators/spherical_harmonics.html), [`o3.Linear`](https://hpc-ai-team.github.io/LunaEquivariance/operators/linear.html), [`o3.Norm`](https://hpc-ai-team.github.io/LunaEquivariance/operators/norm.html) |
+| Activations and regularization | [`nn.Activation`](https://hpc-ai-team.github.io/LunaEquivariance/operators/activation.html), [`nn.Gate`](https://hpc-ai-team.github.io/LunaEquivariance/operators/gate.html), [`nn.NormActivation`](https://hpc-ai-team.github.io/LunaEquivariance/operators/norm_activation.html), [`nn.Dropout`](https://hpc-ai-team.github.io/LunaEquivariance/operators/dropout.html) |
+| Normalization | [`nn.EquivariantNorm`](https://hpc-ai-team.github.io/LunaEquivariance/operators/equivariant_norm.html), [`nn.GraphSoftmax`](https://hpc-ai-team.github.io/LunaEquivariance/operators/graph_softmax.html) |
+| Self-attention | [`nn.VarlenAttention`](https://hpc-ai-team.github.io/LunaEquivariance/operators/varlen_attention.html), [`nn.VarlenAttentionConv`](https://hpc-ai-team.github.io/LunaEquivariance/operators/varlen_attention_conv.html) |
+
+## How to use LunaEquivariance
+
+LunaEquivariance provides module and functional interfaces. See the
+[API reference](https://hpc-ai-team.github.io/LunaEquivariance/api.html) for operator usage and
+[examples](examples/) for training and e3nn integration.
+
+UVU tensor product convolution using the module interface:
+
+```python
+import torch
+from LunaEquivariance import TensorProductConv
+
+dtype = torch.float32  # Use torch.float64 for FP64.
+
+# Define module
+conv = TensorProductConv(
+    "32x1o", "1x1o", "32x0e", [(0, 0, 0, "uvu", True)],
+    internal_weights=False, shared_weights=False, dtype=dtype,
+).cuda()
+
+# Alternatively, initialize from an e3nn tensor product:
+# from e3nn import o3
+# tp = o3.TensorProduct(
+#     "32x1o", "1x1o", "32x0e", [(0, 0, 0, "uvu", True)],
+#     internal_weights=False, shared_weights=False,
+# ).to(device="cuda", dtype=dtype)
+# conv = TensorProductConv(tp)
+
+# Prepare inputs
+num_atoms, num_edges = 4, 8
+src = torch.tensor([1, 3, 0, 2, 1, 3, 0, 2], device="cuda")
+dst = torch.tensor([0, 0, 1, 1, 2, 2, 3, 3], device="cuda")  # Edges sorted by destination.
+graph = conv.create_graph(src, dst, num_atoms)  # Reuse while connectivity is unchanged.
+x = torch.randn(num_atoms, conv.irreps_in1.dim, device="cuda", dtype=dtype, requires_grad=True)
+edge_attr = torch.randn(num_edges, conv.irreps_in2.dim, device="cuda", dtype=dtype)
+weight = torch.randn(num_edges, conv.weight_numel, device="cuda", dtype=dtype, requires_grad=True)
+
+# Forward
+y = conv(x, edge_attr, weight, graph=graph)  # (num_atoms, conv.irreps_out.dim)
+```
+
+Using the functional interface with the same inputs:
+
+```python
+from LunaEquivariance import prepare_tensor_product, tensor_product_conv
+
+# Initialize once
+config = prepare_tensor_product(
+    "32x1o", "1x1o", "32x0e", [(0, 0, 0, "uvu", True)],
+    shared_weights=False, dtype=dtype, convolution=True,
+)
+
+# Forward
+y = tensor_product_conv(x, edge_attr, weight, graph=graph, config=config)
+```
+
+Check your inputs against e3nn/torch, including numerical errors, timings and speedup:
+
+```python
+from LunaEquivariance.utils.testing import check_compatibility
+
+report = check_compatibility(
+    conv, (x, edge_attr, weight), kwargs={"graph": graph}, derivative_order=2,
+)
+```
+
+## Performance
+
+Speedups over e3nn for UVU tensor product convolution on H100, covering forward, backward, double backward and triple backward.
+
+In FP32, LunaTP achieves **58.9–151.8×** average speedup.
+
+![UVU convolution speedup and memory savings over e3nn in FP32](docs/assets/kernel_bench_result_h100.png)
+
+In FP64, LunaTP achieves **45.6–133.7×** average speedup.
+
+![UVU convolution speedup and memory savings over e3nn in FP64](docs/assets/kernel_bench_result_h100_fp64.png)
+
+## Contact
+
+Questions, issues, or feature requests? Feel free to [open an issue](https://github.com/HPC-AI-Team/LunaEquivariance/issues) or [email me](mailto:zhouyuanchang23s@ict.ac.cn).
+
+## Citation
+
+If you use LunaEquivariance in your research, please cite:
+
+```bibtex
+@misc{zhou2026lunatpenablingefficientscalable,
+    title={LunaTP: Enabling Efficient and Scalable Tensor Product for Equivariant MLIPs},
+    author={Yuanchang Zhou and Mingzhen Li and Chen Wang and Xiangyu Zhang and Guangming Tan and Weile Jia},
+    year={2026},
+    eprint={2610.18393},
+    archivePrefix={arXiv},
+    url={https://arxiv.org/abs/2610.18393},
+}
+```
+
+## License
+
+LunaEquivariance is licensed under the [BSD 3-Clause License](LICENSE).