From 5db1bd733201dc8f64ea5309ae76a8d43e024214 Mon Sep 17 00:00:00 2001
From: Linux-cpp-lisp <1473644+Linux-cpp-lisp@users.noreply.github.com>
Date: Sat, 29 Jun 2024 18:04:21 -0400
Subject: [PATCH 1/8] fix version check

---
 nequip/__init__.py | 6 ++----
 1 file changed, 2 insertions(+), 4 deletions(-)

diff --git a/nequip/__init__.py b/nequip/__init__.py
index 654cd78a..537242b3 100644
--- a/nequip/__init__.py
+++ b/nequip/__init__.py
@@ -11,10 +11,8 @@
 torch_version = packaging.version.parse(torch.__version__)
 
 # only allow 1.11*, 1.13* or higher (no 1.12.*)
-assert (torch_version > packaging.version.parse("1.11.0")) and not (
-    packaging.version.parse("1.12.0")
-    <= torch_version
-    < packaging.version.parse("1.13.0")
+assert (torch_version == packaging.version.parse("1.11")) or (
+    torch_version >= packaging.version.parse("1.13")
 ), f"NequIP supports PyTorch 1.11.* or 1.13.* or later, but {torch_version} found"
 
 # warn if using 1.13* or 2.0.*

From 8d5179e0beed83a7d932251df757b78a6c68db4f Mon Sep 17 00:00:00 2001
From: Linux-cpp-lisp <1473644+Linux-cpp-lisp@users.noreply.github.com>
Date: Sat, 29 Jun 2024 18:05:49 -0400
Subject: [PATCH 2/8] Fix performance warning

---
 nequip/__init__.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/nequip/__init__.py b/nequip/__init__.py
index 537242b3..ce145b41 100644
--- a/nequip/__init__.py
+++ b/nequip/__init__.py
@@ -16,9 +16,9 @@
 ), f"NequIP supports PyTorch 1.11.* or 1.13.* or later, but {torch_version} found"
 
 # warn if using 1.13* or 2.0.*
-if packaging.version.parse("1.13.0") <= torch_version < packaging.version.parse("2.1"):
+if packaging.version.parse("1.13.0") <= torch_version:
     warnings.warn(
-        f"!! PyTorch version {torch_version} found. Upstream issues in PyTorch versions 1.13.* and 2.0.* have been seen to cause unusual performance degredations on some CUDA systems that become worse over time; see https://github.com/mir-group/nequip/discussions/311. The best tested PyTorch version to use with CUDA devices is 1.11; while using other versions if you observe this problem, an unexpected lack of this problem, or other strange behavior, please post in the linked GitHub issue."
+        f"!! PyTorch version {torch_version} found. Upstream issues in PyTorch versions 1.13.* and 2.* have been seen to cause unusual performance degredations on some CUDA systems that become worse over time; see https://github.com/mir-group/nequip/discussions/311. The best tested PyTorch version to use with CUDA devices is 1.11; while using other versions if you observe this problem, an unexpected lack of this problem, or other strange behavior, please post in the linked GitHub issue."
     )
 
 

From ee03f0c540f584129507dc4f74d441240a1cbf3c Mon Sep 17 00:00:00 2001
From: Linux-cpp-lisp <1473644+Linux-cpp-lisp@users.noreply.github.com>
Date: Mon, 1 Jul 2024 15:54:46 -0400
Subject: [PATCH 3/8] fewer spurious failures

---
 nequip/utils/unittests/conftest.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/nequip/utils/unittests/conftest.py b/nequip/utils/unittests/conftest.py
index a2dc103d..59bee4ec 100644
--- a/nequip/utils/unittests/conftest.py
+++ b/nequip/utils/unittests/conftest.py
@@ -45,7 +45,7 @@
 # The default float tolerance
 FLOAT_TOLERANCE = {
     t: torch.as_tensor(v, dtype=dtype_from_name(t))
-    for t, v in {"float32": 1e-3, "float64": 1e-10}.items()
+    for t, v in {"float32": 1e-3, "float64": 1e-8}.items()
 }
 
 

From 0eabec7c421f2330adbdc7e9d50c2d649563fc7f Mon Sep 17 00:00:00 2001
From: Linux-cpp-lisp <1473644+Linux-cpp-lisp@users.noreply.github.com>
Date: Mon, 1 Jul 2024 16:07:27 -0400
Subject: [PATCH 4/8] take FLOAT_TOLERANCE from nequip, not e3nn

---
 nequip/utils/test.py               | 22 +++++++++++++++++-----
 nequip/utils/unittests/conftest.py |  9 +--------
 2 files changed, 18 insertions(+), 13 deletions(-)

diff --git a/nequip/utils/test.py b/nequip/utils/test.py
index de843418..4c7450b5 100644
--- a/nequip/utils/test.py
+++ b/nequip/utils/test.py
@@ -2,7 +2,7 @@
 
 import torch
 from e3nn import o3
-from e3nn.util.test import equivariance_error, FLOAT_TOLERANCE
+from e3nn.util.test import equivariance_error
 
 from nequip.nn import GraphModuleMixin, GraphModel
 from nequip.data import (
@@ -12,7 +12,17 @@
     _EDGE_FIELDS,
     _CARTESIAN_TENSOR_FIELDS,
 )
-
+from nequip.utils.misc import dtype_from_name
+
+# The default float tolerance
+FLOAT_TOLERANCE = {
+    t: torch.as_tensor(v, dtype=dtype_from_name(t))
+    for t, v in {"float32": 1e-3, "float64": 1e-10}.items()
+}
+# Allow lookup by name or dtype object:
+for t, v in list(FLOAT_TOLERANCE.items()):
+    FLOAT_TOLERANCE[dtype_from_name(t)] = v
+del t, v
 
 # This has to be somewhat large because of float32 sum reductions over many edges/atoms
 PERMUTATION_FLOAT_TOLERANCE = {torch.float32: 1e-4, torch.float64: 1e-10}
@@ -46,9 +56,11 @@ def assert_permutation_equivariant(
 
     if tolerance is None:
         atol = PERMUTATION_FLOAT_TOLERANCE[
-            func.model_dtype
-            if isinstance(func, GraphModel)
-            else torch.get_default_dtype()
+            (
+                func.model_dtype
+                if isinstance(func, GraphModel)
+                else torch.get_default_dtype()
+            )
         ]
     else:
         atol = tolerance
diff --git a/nequip/utils/unittests/conftest.py b/nequip/utils/unittests/conftest.py
index 59bee4ec..aa716d07 100644
--- a/nequip/utils/unittests/conftest.py
+++ b/nequip/utils/unittests/conftest.py
@@ -12,12 +12,11 @@
 
 import torch
 
-from nequip.utils.test import set_irreps_debug
+from nequip.utils.test import set_irreps_debug, FLOAT_TOLERANCE
 from nequip.data import AtomicData, ASEDataset
 from nequip.data.transforms import TypeMapper
 from nequip.utils.torch_geometric import Batch
 from nequip.utils._global_options import _set_global_options
-from nequip.utils.misc import dtype_from_name
 
 # Sometimes we run parallel using pytest-xdist, and want to be able to use
 # as many GPUs as are available
@@ -42,12 +41,6 @@
     # Test parallelization, but don't waste time spawning tons of workers if lots of cores available
     os.environ["NEQUIP_NUM_TASKS"] = "2"
 
-# The default float tolerance
-FLOAT_TOLERANCE = {
-    t: torch.as_tensor(v, dtype=dtype_from_name(t))
-    for t, v in {"float32": 1e-3, "float64": 1e-8}.items()
-}
-
 
 @pytest.fixture(scope="session", autouse=True, params=["float32", "float64"])
 def float_tolerance(request):

From a4650260cc20bf3c03660b124eb43971f6f10a51 Mon Sep 17 00:00:00 2001
From: Linux-cpp-lisp <1473644+Linux-cpp-lisp@users.noreply.github.com>
Date: Mon, 1 Jul 2024 16:09:33 -0400
Subject: [PATCH 5/8] relax tolerance for F64

---
 nequip/utils/unittests/model_tests.py | 18 +++++++++++++-----
 1 file changed, 13 insertions(+), 5 deletions(-)

diff --git a/nequip/utils/unittests/model_tests.py b/nequip/utils/unittests/model_tests.py
index 37e9dcb6..b9d6790b 100644
--- a/nequip/utils/unittests/model_tests.py
+++ b/nequip/utils/unittests/model_tests.py
@@ -228,7 +228,13 @@ def test_equivariance(self, model, atomic_batch, device):
         instance, out_fields = model
         instance = instance.to(device=device)
         atomic_batch = atomic_batch.to(device=device)
-        assert_AtomicData_equivariant(func=instance, data_in=atomic_batch)
+        assert_AtomicData_equivariant(
+            func=instance,
+            data_in=atomic_batch,
+            e3_tolerance={torch.float32: 1e-3, torch.float64: 1e-8}[
+                torch.get_default_dtype()
+            ],
+        )
 
     def test_embedding_cutoff(self, model, config, device):
         instance, out_fields = model
@@ -449,10 +455,12 @@ def test_partial_forces(self, config, atomic_batch, device, strict_locality):
                 assert torch.allclose(
                     output[k],
                     output_partial[k],
-                    atol=1e-8
-                    if k == AtomicDataDict.TOTAL_ENERGY_KEY
-                    and torch.get_default_dtype() == torch.float64
-                    else 1e-5,
+                    atol=(
+                        1e-8
+                        if k == AtomicDataDict.TOTAL_ENERGY_KEY
+                        and torch.get_default_dtype() == torch.float64
+                        else 1e-5
+                    ),
                 )
             else:
                 assert torch.equal(output[k], output_partial[k])

From b73bf6f32b93cdc6dfbeb77cf79231109e869438 Mon Sep 17 00:00:00 2001
From: cw-tan <chuinwei.tan@gmail.com>
Date: Tue, 9 Jul 2024 10:08:48 -0400
Subject: [PATCH 6/8] prepare for 0.6.1 release

---
 CHANGELOG.md       | 5 ++++-
 nequip/_version.py | 2 +-
 2 files changed, 5 insertions(+), 2 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index bfab0c68..e1c228d0 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,7 +7,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 Most recent change on the bottom.
 
 
-## Unreleased - 0.6.1
+## Unreleased
+
+
+## [0.6.1] - 2024-7-9
 ### Added
 - add support for equivariance testing of arbitrary Cartesian tensor outputs
 - [Breaking] use entry points for `nequip.extension`s (e.g. for field registration)
diff --git a/nequip/_version.py b/nequip/_version.py
index 8e22989a..6c1533d0 100644
--- a/nequip/_version.py
+++ b/nequip/_version.py
@@ -2,4 +2,4 @@
 # See Python packaging guide
 # https://packaging.python.org/guides/single-sourcing-package-version/
 
-__version__ = "0.6.0"
+__version__ = "0.6.1"

From feff269f8f273f041b906c63724dac86d629f762 Mon Sep 17 00:00:00 2001
From: cw-tan <chuinwei.tan@gmail.com>
Date: Tue, 9 Jul 2024 11:09:53 -0400
Subject: [PATCH 7/8] fix torch-numpy dependency in github workflows

---
 .github/workflows/tests.yml | 1 +
 1 file changed, 1 insertion(+)

diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
index 0fb33150..b1e75392 100644
--- a/.github/workflows/tests.yml
+++ b/.github/workflows/tests.yml
@@ -31,6 +31,7 @@ jobs:
       run: |
         python -m pip install --upgrade pip
         pip install setuptools wheel
+        if [ ${TORCH} == 1.13.1 ]; then pip install numpy==1.26.4; fi  # older torch versions fail with numpy 2
         pip install torch==${TORCH} -f https://download.pytorch.org/whl/cpu/torch_stable.html
         pip install h5py scikit-learn  # install packages that aren't required dependencies but that the tests do need
         pip install  --upgrade-strategy only-if-needed .

From 69385abe0e1a2cfed25851336fe17ec1c6571b71 Mon Sep 17 00:00:00 2001
From: Chuin Wei Tan <87742566+cw-tan@users.noreply.github.com>
Date: Tue, 9 Jul 2024 11:34:30 -0400
Subject: [PATCH 8/8] Update github workflows

Co-authored-by: Alby M. <1473644+Linux-cpp-lisp@users.noreply.github.com>
---
 .github/workflows/tests.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
index b1e75392..b09cb076 100644
--- a/.github/workflows/tests.yml
+++ b/.github/workflows/tests.yml
@@ -31,7 +31,7 @@ jobs:
       run: |
         python -m pip install --upgrade pip
         pip install setuptools wheel
-        if [ ${TORCH} == 1.13.1 ]; then pip install numpy==1.26.4; fi  # older torch versions fail with numpy 2
+        if [ ${TORCH} = "1.13.1" ]; then pip install numpy==1.*; fi  # older torch versions fail with numpy 2
         pip install torch==${TORCH} -f https://download.pytorch.org/whl/cpu/torch_stable.html
         pip install h5py scikit-learn  # install packages that aren't required dependencies but that the tests do need
         pip install  --upgrade-strategy only-if-needed .