[JAX] Remove stale references to //jax:experimental, which is being removed.

PiperOrigin-RevId: 860159477

Author: hawkinsp

distrax/_src/distributions/log_stddev_normal_test.py

@@ -22,7 +22,6 @@
 from distrax._src.distributions import normal
 from distrax._src.utils import compat
 import jax
-import jax.experimental
 import jax.numpy as jnp
 import mock
 import numpy as np

distrax/_src/distributions/multinomial_test.py

@@ -23,7 +23,6 @@
 from distrax._src.utils import equivalence
 from distrax._src.utils import math
 import jax
-import jax.experimental
 import jax.numpy as jnp
 import numpy as np
 from scipy import stats

distrax/_src/distributions/mvn_diag_plus_low_rank_test.py

@@ -23,7 +23,6 @@
 from distrax._src.utils import equivalence
 
 import jax
-import jax.experimental
 import jax.numpy as jnp
 import numpy as np
 from tensorflow_probability.substrates import jax as tfp

distrax/_src/distributions/mvn_diag_test.py

@@ -23,7 +23,6 @@
 from distrax._src.utils import compat
 from distrax._src.utils import equivalence
 import jax
-import jax.experimental
 import jax.numpy as jnp
 import numpy as np
 

distrax/_src/distributions/one_hot_categorical_test.py

@@ -24,7 +24,6 @@
 from distrax._src.utils import equivalence
 from distrax._src.utils import math
 import jax
-import jax.experimental
 import jax.numpy as jnp
 import numpy as np
 import scipy