The version of JAX has been restricted to between 0.4.1 and 0.4.24. &…

… Fix dtype bug in NoCutoffNBList
deepmodeling · Mar 7, 2024 · e299ced · e299ced
1 parent 285dd85
commit e299ced
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 1 deletion.
diff --git a/dmff/common/nblist.py b/dmff/common/nblist.py
@@ -88,7 +88,7 @@ def _do_cov_map(self, pairs):
     def allocate(self, coords, box=None):
         self._positions = coords  # cache it
         natoms = coords.shape[0]
-        nblist = np.fromiter(permutations(range(natoms), 2), dtype=np.dtype(int, 2))
+        nblist = np.fromiter(permutations(range(natoms), 2), dtype=np.dtype((int, 2)))
         nlist = nblist[nblist[:, 0] < nblist[:, 1]]
         if self.capacity_multiplier is None:
             self.capacity_multiplier = int(nlist.shape[0] * 1.3)

diff --git a/setup.py b/setup.py
@@ -20,6 +20,7 @@
 
 install_requires = [
     "numpy>=1.18",
+    "jax<=0.4.24, >=0.4.1",
     "openmm>=7.6.0",
     "freud-analysis",
     "networkx>=3.0",