Add IrDictPolynomial and _ir_dict descriptor variants

cuequivariance/cuequivariance/__init__.py

@@ -12,6 +12,10 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+"""cuEquivariance - GPU-accelerated equivariant operations for 3D neural networks.
+
+For AI coding assistants: run `python -m cuequivariance skill` for detailed usage guidance.
+"""
 import importlib.resources
 
 __version__ = (
@@ -51,6 +55,8 @@
     reduced_antisymmetric_tensor_product_basis,
     EquivariantPolynomial,
     EquivariantTensorProduct,  # deprecated
+    IrDictPolynomial,
+    split_polynomial_by_irreps,
 )
 
 from cuequivariance import segmented_polynomials as segmented_polynomials
@@ -89,6 +95,8 @@
     "reduced_antisymmetric_tensor_product_basis",
     "EquivariantPolynomial",
     "EquivariantTensorProduct",
+    "IrDictPolynomial",
+    "split_polynomial_by_irreps",
     "segmented_polynomials",
     "group_theory",
     "descriptors",

cuequivariance/cuequivariance/__main__.py

@@ -0,0 +1,29 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import sys
+from pathlib import Path
+
+
+def main():
+    if len(sys.argv) >= 2 and sys.argv[1] == "skill":
+        skill_path = Path(__file__).parent / "SKILL.md"
+        print(skill_path.read_text())
+    else:
+        print("Usage: python -m cuequivariance skill")
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()

cuequivariance/cuequivariance/etc/linalg.py

@@ -107,6 +107,10 @@ def limit_denominator(n, d, max_denominator: int):
     n1, d1 = p0 + k * p1, q0 + k * q1
     n2, d2 = p1, q1
     with np.errstate(over="ignore"):
+        # The intermediate products (n2*d0, n0*d2) overflow int64 (~2^102), but the
+        # overflow is benign: their difference is bounded by d0 < 2^62 (fits in int64),
+        # and two's complement subtraction recovers it exactly. The final product
+        # d1*(difference) is also bounded by d0 < 2^63.
         mask = np.abs(d1 * (n2 * d0 - n0 * d2)) <= np.abs(d2 * (n1 * d0 - n0 * d1))
     return np.where(
         d0 <= max_denominator,

cuequivariance/cuequivariance/group_theory/__init__.py

@@ -44,6 +44,7 @@
 
 from .equivariant_polynomial import EquivariantPolynomial
 from .equivariant_tensor_product import EquivariantTensorProduct
+from .ir_dict_polynomial import IrDictPolynomial, split_polynomial_by_irreps
 
 
 __all__ = [
@@ -72,4 +73,6 @@
     "reduced_antisymmetric_tensor_product_basis",
     "EquivariantPolynomial",
     "EquivariantTensorProduct",
+    "IrDictPolynomial",
+    "split_polynomial_by_irreps",
 ]

cuequivariance/cuequivariance/group_theory/descriptors/__init__.py

@@ -15,12 +15,17 @@
 from .transposition import transpose
 from .irreps_tp import (
     full_tensor_product,
+    full_tensor_product_ir_dict,
     fully_connected_tensor_product,
+    fully_connected_tensor_product_ir_dict,
     channelwise_tensor_product,
+    channelwise_tensor_product_ir_dict,
     elementwise_tensor_product,
+    elementwise_tensor_product_ir_dict,
     linear,
+    linear_ir_dict,
 )
-from .symmetric_contractions import symmetric_contraction
+from .symmetric_contractions import symmetric_contraction, symmetric_contraction_ir_dict
 from .rotations import (
     fixed_axis_angle_rotation,
     y_rotation,
@@ -30,16 +35,26 @@
     yxy_rotation,
     inversion,
 )
-from .spherical_harmonics_ import sympy_spherical_harmonics, spherical_harmonics
+from .spherical_harmonics_ import (
+    sympy_spherical_harmonics,
+    spherical_harmonics,
+    spherical_harmonics_ir_dict,
+)
 
 __all__ = [
     "transpose",
     "full_tensor_product",
+    "full_tensor_product_ir_dict",
     "fully_connected_tensor_product",
+    "fully_connected_tensor_product_ir_dict",
     "channelwise_tensor_product",
+    "channelwise_tensor_product_ir_dict",
     "elementwise_tensor_product",
+    "elementwise_tensor_product_ir_dict",
     "linear",
+    "linear_ir_dict",
     "symmetric_contraction",
+    "symmetric_contraction_ir_dict",
     "fixed_axis_angle_rotation",
     "y_rotation",
     "x_rotation",
@@ -49,4 +64,5 @@
     "inversion",
     "sympy_spherical_harmonics",
     "spherical_harmonics",
+    "spherical_harmonics_ir_dict",
 ]