Merge branch 'main' into fix-shared-weights

.github/workflows/doc_build.yml

@@ -2,7 +2,7 @@ name: Build Documentation
 
 on:
   pull_request:
-    branches: [ "main" ]
+    branches: [ "main", "release" ]
 
 jobs:
   build:
@@ -17,11 +17,12 @@ jobs:
       - name: Install dependencies
         run: |
           python -m pip install --upgrade pip
-          python -m pip install pytest
-          python -m pip install ./cuequivariance
-          python -m pip install ./cuequivariance_jax
-          python -m pip install ./cuequivariance_torch
-          pip install -r docs/requirements.txt
+          python -m pip install --upgrade uv
+          python -m uv pip install pytest
+          python -m uv pip install ./cuequivariance
+          python -m uv pip install ./cuequivariance_jax
+          python -m uv pip install ./cuequivariance_torch
+          python -m uv pip install -r docs/requirements.txt
       - name: Build sphinx
         run: |
           sphinx-build -b html docs docs/public

.github/workflows/doc_build_deploy.yml

@@ -3,7 +3,7 @@ name: Build Documentation and Deploy
 on:
   # Runs on pushes targeting the default branch
   push:
-    branches: ["main"]
+    branches: [ "main", "release" ]
 
   # Allows you to run this workflow manually from the Actions tab
   workflow_dispatch:
@@ -38,11 +38,12 @@ jobs:
       - name: Install dependencies
         run: |
           python -m pip install --upgrade pip
-          python -m pip install pytest
-          python -m pip install ./cuequivariance
-          python -m pip install ./cuequivariance_jax
-          python -m pip install ./cuequivariance_torch
-          pip install -r docs/requirements.txt
+          python -m pip install --upgrade uv
+          python -m uv pip install pytest
+          python -m uv pip install ./cuequivariance
+          python -m uv pip install ./cuequivariance_jax
+          python -m uv pip install ./cuequivariance_torch
+          python -m uv pip install -r docs/requirements.txt
       - name: Build sphinx
         run: |
           sphinx-build -b html docs docs/public

.github/workflows/style.yml

@@ -2,9 +2,9 @@ name: Run Style Check
 
 on:
   push:
-    branches: [ "main" ]
+    branches: [ "main", "release" ]
   pull_request:
-    branches: [ "main" ]
+    branches: [ "main", "release" ]
 
 
 jobs:
@@ -23,7 +23,8 @@ jobs:
     - name: Install dependencies
       run: |
         python -m pip install --upgrade pip
-        python -m pip install ruff==0.8.2
+        python -m pip install --upgrade uv
+        python -m uv pip install ruff==0.8.2
     - name: Style check
       run: |
         ruff check cuequivariance

.github/workflows/tests.yml

@@ -2,9 +2,9 @@ name: Run Tests
 
 on:
   push:
-    branches: [ "main" ]
+    branches: [ "main", "release" ]
   pull_request:
-    branches: [ "main" ]
+    branches: [ "main", "release" ]
 
 
 jobs:
@@ -25,8 +25,9 @@ jobs:
     - name: Install dependencies
       run: |
         python -m pip install --upgrade pip
-        python -m pip install pytest
-        python -m pip install ./cuequivariance
+        python -m pip install --upgrade uv
+        python -m uv pip install pytest
+        python -m uv pip install ./cuequivariance
     - name: Test with pytest
       run: |
         pytest --doctest-modules -x cuequivariance
@@ -48,9 +49,10 @@ jobs:
     - name: Install dependencies
       run: |
         python -m pip install --upgrade pip
-        python -m pip install pytest
-        python -m pip install ./cuequivariance
-        python -m pip install ./cuequivariance_jax
+        python -m pip install --upgrade uv
+        python -m uv pip install pytest
+        python -m uv pip install ./cuequivariance
+        python -m uv pip install ./cuequivariance_jax
     - name: Test with pytest
       run: |
         pytest --doctest-modules -x cuequivariance_jax
@@ -78,9 +80,10 @@ jobs:
     - name: Install dependencies
       run: |
         python -m pip install --upgrade pip
-        python -m pip install pytest torch e3nn
-        python -m pip install ./cuequivariance
-        python -m pip install ./cuequivariance_torch
+        python -m pip install --upgrade uv
+        python -m uv pip install pytest torch e3nn
+        python -m uv pip install ./cuequivariance
+        python -m uv pip install ./cuequivariance_torch
     - name: Test with pytest
       run: |
         pytest --doctest-modules -x cuequivariance_torch

CHANGELOG.md

@@ -11,7 +11,7 @@
 
 ### Added
 
-- Support of `torch.jit.script` and `torch.compile`. Known issue: the export in c++ is not working.
+- Add an experimental support for `torch.jit.script` and `torch.compile`. Known issue: the export in c++ is not working.
 - Add `cue.IrrepsAndLayout`: A simple class that inherits from `cue.Rep` and contains a `cue.Irreps` and a `cue.IrrepsLayout`.
 - Add `cuex.RepArray` for representing an array of any kind of representations (not only irreps like before with `cuex.IrrepsArray`).
 

LICENSES/Third_party_attr.txt

@@ -1,5 +1,5 @@
 Name: Jinja2
- Version: 3.1.4
+ Version: 3.1.5
  License: BSD License 
  URL: https://github.com/pallets/jinja/
  License Text:
@@ -103,7 +103,7 @@ For more information, please refer to <http://unlicense.org>
 *******************
 
 Name: fsspec
- Version: 2024.10.0
+ Version: 2024.12.0
  License: BSD License 
  URL: https://github.com/fsspec/filesystem_spec
  License Text:
@@ -167,7 +167,7 @@ Name: iniconfig
 *******************
 
 Name: jax
- Version: 0.4.35
+ Version: 0.4.38
  License: Apache-2.0 
  URL: https://github.com/jax-ml/jax
  License Text:
@@ -377,7 +377,7 @@ Name: jax
 *******************
 
 Name: jaxlib
- Version: 0.4.35
+ Version: 0.4.38
  License: Apache-2.0 
  URL: https://github.com/jax-ml/jax
  License Text:
@@ -4891,7 +4891,7 @@ freely, subject to the following restrictions:
 *******************
 
 Name: ml_dtypes
- Version: 0.5.0
+ Version: 0.5.1
  License: Apache Software License 
  URL: https://github.com/jax-ml/ml_dtypes
  License Text:
@@ -5181,7 +5181,7 @@ NetworkX is distributed with the 3-clause BSD license.
 *******************
 
 Name: numpy
- Version: 2.1.3
+ Version: 2.2.1
  License: BSD License 
  URL: https://numpy.org
  License Text:
@@ -6201,7 +6201,7 @@ under the terms of *both* these licenses.
 *******************
 
 Name: pip
- Version: 24.1.1
+ Version: 24.3.1
  License: MIT License 
  URL: https://pip.pypa.io/
  License Text:
@@ -6258,7 +6258,7 @@ SOFTWARE.
 *******************
 
 Name: pytest
- Version: 8.3.3
+ Version: 8.3.4
  License: MIT License 
  URL: https://docs.pytest.org/en/latest/
  License Text:
@@ -6287,7 +6287,7 @@ SOFTWARE.
 *******************
 
 Name: scipy
- Version: 1.14.1
+ Version: 1.15.1
  License: BSD License 
  URL: https://scipy.org/
  License Text:
@@ -6324,11 +6324,12 @@ OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
 ----
 
-This binary distribution of SciPy also bundles the following software:
+This binary distribution of SciPy can also bundle the following software
+(depending on the build):
 
 
 Name: OpenBLAS
-Files: scipy.libs/libopenblas*.so
+Files: scipy.libs/libscipy_openblas*.so
 Description: bundled as a dynamically linked library
 Availability: https://github.com/OpenMathLib/OpenBLAS/
 License: BSD-3-Clause-Attribution
@@ -6364,7 +6365,7 @@ License: BSD-3-Clause-Attribution
 
 
 Name: LAPACK
-Files: scipy.libs/libopenblas*.so
+Files: scipy.libs/libscipy_openblas*.so
 Description: bundled in OpenBLAS
 Availability: https://github.com/OpenMathLib/OpenBLAS/
 License: BSD-3-Clause-Attribution
@@ -7228,7 +7229,7 @@ License: LGPL-2.1-or-later
 *******************
 
 Name: setuptools
- Version: 75.5.0
+ Version: 75.8.0
  License: MIT License 
  URL: https://github.com/pypa/setuptools
  License Text:
@@ -7414,7 +7415,7 @@ SOFTWARE.
 *******************
 
 Name: tomli
- Version: 2.1.0
+ Version: 2.2.1
  License: MIT License 
  URL: https://github.com/hukkin/tomli
  License Text:

cuequivariance_jax/cuequivariance_jax/rep_array/jax_rep_array.py

@@ -1,3 +1,17 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 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.
 from __future__ import annotations
 
 from dataclasses import dataclass, field

cuequivariance_torch/tests/__init__.py

@@ -0,0 +1,14 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 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.