diff --git a/doc/examples/example_jax_petab/ExampleJaxPEtab.ipynb b/doc/examples/example_jax_petab/ExampleJaxPEtab.ipynb
index 5ade5cf528..30feb2db9b 100644
--- a/doc/examples/example_jax_petab/ExampleJaxPEtab.ipynb
+++ b/doc/examples/example_jax_petab/ExampleJaxPEtab.ipynb
@@ -465,7 +465,7 @@
     "# initialise adam\n",
     "optim = adam(0.01)\n",
     "# eqx.partition is necessary here to only initialize the optimizer for array variables\n",
-    "param, static = eqx.partition(jax_problem, eqx.is_array)\n",
+    "param, static = eqx.partition(jax_problem, eqx.is_inexact_array)\n",
     "opt_state = optim.init(param)\n",
     "\n",
     "\n",
diff --git a/doc/rtd_requirements.txt b/doc/rtd_requirements.txt
index 4211b95cc9..2c1e4784a6 100644
--- a/doc/rtd_requirements.txt
+++ b/doc/rtd_requirements.txt
@@ -9,7 +9,7 @@ setuptools>=67.7.2
 git+https://github.com/jmuhlich/pysb@22d69a350b472f33d85ba64ffb10b190483c1c98
 # For forward type definition in generate_equinox
 matplotlib>=3.7.1
-optax==0.2.6
+optax
 nbsphinx
 nbformat
 myst-parser
diff --git a/scripts/installAmiciSource.sh b/scripts/installAmiciSource.sh
index 436ac86f97..4c416a5788 100755
--- a/scripts/installAmiciSource.sh
+++ b/scripts/installAmiciSource.sh
@@ -39,7 +39,7 @@ python -m pip install --upgrade pip wheel
 python -m pip install --upgrade pip setuptools cmake_build_extension==0.6.0 numpy petab swig
 python -m pip install git+https://github.com/pysb/pysb@master # for SPM with compartments
 python -m pip install git+https://github.com/patrick-kidger/diffrax@main # for events with direction
-python -m pip install 'optax<0.2.7' # for jax petab notebook
+python -m pip install optax # for jax petab notebook
 AMICI_BUILD_TEMP="${AMICI_PATH}/python/sdist/build/temp" \
   python -m pip install --verbose -e "${AMICI_PATH}/python/sdist[petab,test,vis,jax]" --no-build-isolation
 deactivate