bump the jax version to fix #23

cagrikymk · Jun 25, 2024 · 05bfdd2 · 05bfdd2
1 parent d757feb
commit 05bfdd2
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/jaxreaxff/optimizer.py b/jaxreaxff/optimizer.py
@@ -513,7 +513,7 @@ def add_noise_to_params(params, bounds, scale=0.001):
   # produce noise from [-1,1]
   noise = onp.random.uniform(low = -1.0, high = 1.0,size=len(params))
   new_params = params + (noise * scale) * (bounds[:,1] - bounds[:,0])
-  new_params = jnp.clip(new_params,a_min=bounds[:,0],a_max=bounds[:,1])
+  new_params = jnp.clip(new_params, bounds[:,0], bounds[:,1])
 
   return new_params
 

diff --git a/setup.py b/setup.py
@@ -33,8 +33,8 @@ def get_cuda_version():
 '''
 
 INSTALL_REQUIRES = [
-  'jax>=0.4.13,<=0.4.23',
-  'jaxlib>=0.4.13,<=0.4.23',
+  'jax>=0.4.26, <=0.4.30',
+  'jaxlib>=0.4.26,<=0.4.30',
   'jax_md @ git+https://github.com/cagrikymk/jax-md.git@reaxff_dev#egg=jax_md',
   'numba>=0.56',
   'numpy',