Merge branch 'release' into polynomial

NVIDIA · Mar 1, 2025 · bd5c605 · bd5c605
2 parents 9c9a077 + f5a67a5
commit bd5c605
Show file tree

Hide file tree

Showing 29 changed files with 86 additions and 463 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -1,5 +1,7 @@
 ## Latest Changes
 
+## 0.3.0-rc1
+
 ### Breaking Changes
 - In `cuex.equivariant_tensor_product`, the arguments `dtype_math` and `dtype_output` are renamed `math_dtype` and `output_dtype` respectively. Adding consistency with the rest of the library.
 - In `cuex.equivariant_tensor_product`, the arguments `algorithm`, `precision`, `use_custom_primitive` and `use_custom_kernels` are removed. This is to avoid a proliferation of arguments that are not used in all the implementations. An argument `impl: str` is added instead to select the implementation.
@@ -16,6 +18,7 @@
 ### Added
 - JAX Bindings to the uniform 1d JIT kernel. This kernel handles any kind of non-homogeneous polynomials as long as the contraction pattern (subscripts) have only one index. It handles batched/shared/indexed input/output. The indexed input/output are handled by atomic operations.
 - Added `__mul__` to `cue.EquivariantTensorProduct` to allow rescaling the equivariant tensor product.
+- Added a uniform 1d kernel with scatter/gather fusion under `cuet.primitives.tensor_product.TensorProductUniform4x1dIndexed` and `cuet.primitives.tensor_product.TensorProductUniform3x1dIndexed`.
 
 
 ## 0.2.0 (2025-01-24)

diff --git a/LICENSES/Third_party_attr.txt b/LICENSES/Third_party_attr.txt
diff --git a/README.md b/README.md
@@ -17,8 +17,8 @@ pip install cuequivariance-torch
 pip install cuequivariance  # Installs only the core non-ML components
 
 # CUDA kernels for different CUDA versions
-pip install cuequivariance-ops-torch-cu11
-pip install cuequivariance-ops-torch-cu12
+pip install cuequivariance-ops-jax-cu12
+pip install cuequivariance-ops-torch-cu12  # or cu11
 ```
 
 ## License

diff --git a/VERSION b/VERSION
@@ -1 +1 @@
-0.2.0
+0.3.0rc4
diff --git a/cuequivariance/cuequivariance/__init__.py b/cuequivariance/cuequivariance/__init__.py
@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");

diff --git a/cuequivariance/cuequivariance/equivariant_tensor_product.py b/cuequivariance/cuequivariance/equivariant_tensor_product.py
@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");

diff --git a/cuequivariance/cuequivariance/misc/linalg.py b/cuequivariance/cuequivariance/misc/linalg.py
@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");

diff --git a/cuequivariance/cuequivariance/operation.py b/cuequivariance/cuequivariance/operation.py
@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");

diff --git a/cuequivariance/cuequivariance/segmented_tensor_product/dimensions_dict.py b/cuequivariance/cuequivariance/segmented_tensor_product/dimensions_dict.py
@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");

diff --git a/cuequivariance/pyproject.toml b/cuequivariance/pyproject.toml
@@ -42,4 +42,4 @@ classifiers = [
 
 [tool.hatch.version]
 path = "cuequivariance/VERSION"
-pattern = "(?P<version>\\d+\\.\\d+\\.\\d+)"
+pattern = "(?P<version>\\d+\\.\\d+\\.\\d+(?:[a-z]+\\d+)?)"
diff --git a/cuequivariance/tests/segmented_tensor_product/descriptor_test.py b/cuequivariance/tests/segmented_tensor_product/descriptor_test.py
@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");

diff --git a/cuequivariance_jax/cuequivariance_jax/__init__.py b/cuequivariance_jax/cuequivariance_jax/__init__.py
@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");

diff --git a/cuequivariance_jax/cuequivariance_jax/flax_linen/__init__.py b/cuequivariance_jax/cuequivariance_jax/flax_linen/__init__.py
@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");

diff --git a/cuequivariance_jax/cuequivariance_jax/operations/spherical_harmonics.py b/cuequivariance_jax/cuequivariance_jax/operations/spherical_harmonics.py
@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");

diff --git a/cuequivariance_jax/cuequivariance_jax/primitives/primitives_utils.py b/cuequivariance_jax/cuequivariance_jax/primitives/primitives_utils.py
@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");

diff --git a/cuequivariance_jax/cuequivariance_jax/primitives/segmented_polynomial.py b/cuequivariance_jax/cuequivariance_jax/primitives/segmented_polynomial.py
@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -389,7 +389,9 @@ def segmented_polynomial_jvp(
         jvp_buffer_index,
         polynomial.jvp([not isinstance(t, ad.Zero) for t in tangents]),
         math_dtype,
-        name + "_jvp",
+        name
+        + "_jvp"
+        + "".join("0" if isinstance(t, ad.Zero) else "1" for t in tangents),
         impl=impl,
     )
 

diff --git a/cuequivariance_jax/cuequivariance_jax/primitives/segmented_polynomial_ops_impl.py b/cuequivariance_jax/cuequivariance_jax/primitives/segmented_polynomial_ops_impl.py
@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -25,7 +25,10 @@
 
 
 def sanitize_string(s):
-    return re.sub(r"[^A-Za-z_]", "", s)
+    s = re.sub(r"[^A-Za-z0-9_]", "", s)
+    if s == "" or s[0].isdigit():
+        s = "_" + s
+    return s
 
 
 def segmented_polynomial_ops_impl(

diff --git a/cuequivariance_jax/cuequivariance_jax/primitives/segmented_polynomial_vanilla_impl.py b/cuequivariance_jax/cuequivariance_jax/primitives/segmented_polynomial_vanilla_impl.py
@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");

diff --git a/cuequivariance_jax/pyproject.toml b/cuequivariance_jax/pyproject.toml
@@ -39,4 +39,4 @@ classifiers = [
 
 [tool.hatch.version]
 path = "cuequivariance_jax/VERSION"
-pattern = "(?P<version>\\d+\\.\\d+\\.\\d+)"
+pattern = "(?P<version>\\d+\\.\\d+\\.\\d+(?:[a-z]+\\d+)?)"
diff --git a/cuequivariance_jax/tests/operations/spherical_harmonics_test.py b/cuequivariance_jax/tests/operations/spherical_harmonics_test.py
@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");

diff --git a/cuequivariance_jax/tests/primitives/equivariant_polynomial_test.py b/cuequivariance_jax/tests/primitives/equivariant_polynomial_test.py
@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");

diff --git a/cuequivariance_jax/tests/primitives/segmented_polynomial_test.py b/cuequivariance_jax/tests/primitives/segmented_polynomial_test.py
@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");

diff --git a/cuequivariance_torch/cuequivariance_torch/primitives/transpose.py b/cuequivariance_torch/cuequivariance_torch/primitives/transpose.py
@@ -89,11 +89,6 @@ def __init__(
     ):
         super().__init__()
 
-        if use_fallback is not True:
-            # There is a bug in the CUDA kernel for TransposeIrrepsLayout
-            # TODO remove this when the bug is fixed
-            use_fallback = True
-
         info = _transpose_info(segments, device=device)
         self.f = None
 

diff --git a/cuequivariance_torch/pyproject.toml b/cuequivariance_torch/pyproject.toml
@@ -38,4 +38,4 @@ classifiers = [
 
 [tool.hatch.version]
 path = "cuequivariance_torch/VERSION"
-pattern = "(?P<version>\\d+\\.\\d+\\.\\d+)"
+pattern = "(?P<version>\\d+\\.\\d+\\.\\d+(?:[a-z]+\\d+)?)"
diff --git a/docs/api/cuequivariance.rst b/docs/api/cuequivariance.rst
@@ -1,4 +1,4 @@
-.. SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+.. SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
    SPDX-License-Identifier: Apache-2.0
 
    Licensed under the Apache License, Version 2.0 (the "License");

diff --git a/docs/api/cuequivariance_jax.rst b/docs/api/cuequivariance_jax.rst
@@ -1,4 +1,4 @@
-.. SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+.. SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
    SPDX-License-Identifier: Apache-2.0
 
    Licensed under the Apache License, Version 2.0 (the "License");

diff --git a/docs/index.rst b/docs/index.rst
@@ -1,4 +1,4 @@
-.. SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+.. SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
    SPDX-License-Identifier: Apache-2.0
 
    Licensed under the Apache License, Version 2.0 (the "License");

diff --git a/docs/tutorials/beta.rst b/docs/tutorials/beta.rst
@@ -1,4 +1,4 @@
-.. SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+.. SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
    SPDX-License-Identifier: Apache-2.0
 
    Licensed under the Apache License, Version 2.0 (the "License");

diff --git a/docs/tutorials/index.rst b/docs/tutorials/index.rst
@@ -1,4 +1,4 @@
-.. SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+.. SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
    SPDX-License-Identifier: Apache-2.0
 
    Licensed under the Apache License, Version 2.0 (the "License");