update pytree filtering to inexact array (#3143)

Author: BSnelling

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",

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

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