Update requirements

ziatdinovmax · Sep 26, 2024 · e54acc5 · e54acc5
1 parent c81606e
commit e54acc5
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/setup.py b/setup.py
@@ -25,7 +25,8 @@
         packages=find_packages(),
         zip_safe=False,
         install_requires=[
-            'jax>=0.4.0',
+            'jax>=0.4.0,<=0.4.31',
+            'jaxlib>=0.4.0,<=0.4.31',
             'numpyro>=0.13.0',
             'flax>=0.8.4'
         ],
@@ -35,4 +36,4 @@
                      'Operating System :: POSIX :: Linux',
                      'Operating System :: MacOS :: MacOS X',
                      'Topic :: Scientific/Engineering']
-    )
+    )