fix jax version instead of scvi-tools

Author: vitkl

setup.cfg

@@ -15,9 +15,10 @@ include_package_data = True
 packages = find:
 install_requires =
     pyro-ppl>=1.8.0
-    scvi-tools>=0.15.3<0.15.5
+    scvi-tools>=0.15.3
     torch>=1.9.0
     pymc3>=3.8<3.10
+    jax>=0.3,<=0.3.6
     arviz==0.10.0
     numpy
     pandas