Pin JAX version to use old CPU runtime (#570)

CHANGELOG.md

@@ -1,5 +1,9 @@
 # 0.6.0
 
+### Pin of JAX version
+
+Installing `Jaxley` will no longer install the newest version of `JAX`. We realized that, on CPU, with version `jax==0.4.32` or newer, simulation time in `Jaxley` is 10x slower and gradient time is 5x slower as compared to older versions of JAX. Newer versions of `JAX` can be made equally fast as older versions by setting `os.environ['XLA_FLAGS'] = '--xla_cpu_use_thunk_runtime=false'` at the beginning of your jupyter notebook (#570, @michaeldeistler).
+
 ### New Features
 
 - Add ability to record synaptic currents (#523, @ntolley). Recordings can be turned on with

pyproject.toml

@@ -26,7 +26,11 @@ requires-python = ">=3.10"
 readme = "README.md"
 keywords = ["neuroscience", "biophysics", "simulator", "jax"]
 dependencies = [
-    "jax",
+    # See #571 for the pin on the JAX version. Jaxley runs with newer versions of JAX
+    # but is much slower on CPU. Newer versions can be made equally fast by setting
+    # `os.environ['XLA_FLAGS'] = '--xla_cpu_use_thunk_runtime=false'` at the beginning
+    # of a jupyter notebook.
+    "jax <= 0.4.31",
     "matplotlib",
     "numpy",
     "pandas>=2.2.0",
@@ -77,4 +81,4 @@ profile = "black"
 [tool.setuptools.packages.find]
 where = ["."]  # list of folders that contain the packages (["."] by default)
 include = ["*"]  # package names should match these glob patterns (["*"] by default)
-exclude = []  # exclude packages matching these glob patterns (empty by default)
+exclude = []  # exclude packages matching these glob patterns (empty by default)