fix(deps): use jnp.moveaxis instead of jax.batching.moveaxis (#72)

[Jax 0.7.1](https://docs.jax.dev/en/latest/changelog.html) deprecates
`jax.batching.moveaxis` and recommends switching to `jnp.moveaxis`

#### Description of changes
Followed the recommendation!

#### Testing done
CI passes with jax 0.7.1 (and whatever version the lockfile specifies on
github)

Author: jpbrodrick89

tesseract_jax/primitive.py

@@ -5,6 +5,7 @@
 from collections.abc import Sequence
 from typing import Any, TypeVar
 
+import jax.numpy as jnp
 import jax.tree
 import numpy as np
 from jax import ShapeDtypeStruct, dtypes, extend
@@ -246,7 +247,7 @@ def tesseract_dispatch_batching(
 ) -> Any:
     """Defines how to dispatch batch operations such as vmap (which is used by jax.jacobian)."""
     new_args = [
-        arg if ax is batching.not_mapped else batching.moveaxis(arg, ax, 0)
+        arg if ax is batching.not_mapped else jnp.moveaxis(arg, ax, 0)
         for arg, ax in zip(array_args, axes, strict=True)
     ]