⬆️ Update jax requirement

Updates the requirements on [jax](https://github.com/google/jax) to permit the latest version.
- [Release notes](https://github.com/google/jax/releases)
- [Changelog](https://github.com/google/jax/blob/main/CHANGELOG.md)
- [Commits](jax-ml/jax@jax-v0.4.1...jax-v0.4.16)

---
updated-dependencies:
- dependency-name: jax
  dependency-type: direct:production
...

Signed-off-by: dependabot[bot] <[email protected]>

requirements-dev.txt

@@ -1,7 +1,7 @@
 numpy
 numba
 brainpylib
-jax>=0.4.1, <0.4.16
+jax>=0.4.1, <0.4.17
 jaxlib>=0.4.1, <0.4.16
 matplotlib>=3.4
 msgpack

requirements-doc.txt

@@ -2,7 +2,7 @@ numpy
 tqdm
 msgpack
 numba
-jax>=0.4.1, <0.4.16
+jax>=0.4.1, <0.4.17
 jaxlib>=0.4.1, <0.4.16
 matplotlib>=3.4
 scipy>=1.1.0

requirements.txt

@@ -1,5 +1,5 @@
 numpy
-jax>=0.4.1, <0.4.16
+jax>=0.4.1, <0.4.17
 tqdm
 msgpack
 numba

setup.py

@@ -51,7 +51,7 @@
   author_email='[email protected]',
   packages=packages,
   python_requires='>=3.8',
-  install_requires=['numpy>=1.15', 'jax>=0.4.1, <0.4.16', 'tqdm', 'msgpack', 'numba'],
+  install_requires=['numpy>=1.15', 'jax>=0.4.1, <0.4.17', 'tqdm', 'msgpack', 'numba'],
   url='https://github.com/brainpy/BrainPy',
   project_urls={
     "Bug Tracker": "https://github.com/brainpy/BrainPy/issues",