Bump jax from 0.7.2 to 0.8.1 (#777)

* Bump jax from 0.7.2 to 0.8.1

Bumps [jax](https://github.com/jax-ml/jax) from 0.7.2 to 0.8.1.
- [Release notes](https://github.com/jax-ml/jax/releases)
- [Changelog](https://github.com/jax-ml/jax/blob/jax-v0.8.1/CHANGELOG.md)
- [Commits](jax-ml/jax@jax-v0.7.2...jax-v0.8.1)

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

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

* update import path

---------

Signed-off-by: dependabot[bot] <support@github.com>
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
Co-authored-by: Dion Häfner <dion.haefner@simulation.science>

Author: dependabot[bot]

requirements_jax.txt

@@ -1 +1 @@
-jax==0.7.2
+jax==0.8.2

veros/core/special/tdma_.py

@@ -23,7 +23,7 @@
 import jaxlib.mlir.ir as ir
 from jaxlib.mlir.dialects import mhlo
 
-from jax.interpreters.mlir import custom_call
+from jax._src.interpreters.mlir import custom_call
 
 
 if HAS_CPU_EXT: