Bump jax from 0.3.25 to 0.4.6 in /docker/openproblems-python-pytorch (#846)

* Bump jax from 0.3.25 to 0.4.6 in /docker/openproblems-python-pytorch

Bumps [jax](https://github.com/google/jax) from 0.3.25 to 0.4.6.
- [Release notes](https://github.com/google/jax/releases)
- [Changelog](https://github.com/google/jax/blob/main/CHANGELOG.md)
- [Commits](jax-ml/jax@jax-v0.3.25...jax-v0.4.6)

---
updated-dependencies:
- dependency-name: jax
  dependency-type: direct:production
  update-type: version-update:semver-minor
...

Signed-off-by: dependabot[bot] <support@github.com>

* Bump jaxlib

---------

Signed-off-by: dependabot[bot] <support@github.com>
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
Co-authored-by: Scott Gigante <84813314+scottgigante-immunai@users.noreply.github.com>

Author: dependabot[bot]

docker/openproblems-python-pytorch/requirements.txt

@@ -1,8 +1,8 @@
 git+https://github.com/BayraktarLab/cell2location.git@47c8d6dc90dd3f1ab639861e8617c6ef0b62bb89
 git+https://github.com/czbiohub/molecular-cross-validation@04d9df0
 git+https://github.com/michalk8/neuralee@8946abf  # contains gradient error fix
-jax==0.3.25
-jaxlib==0.3.25
+jax==0.4.6
+jaxlib==0.4.6
 pymde==0.1.18
 scalex==1.0.2
 scikit-misc==0.1.*