Update references to JAX's GitHub repo

JAX has moved from https://github.com/google/jax to https://github.com/jax-ml/jax

PiperOrigin-RevId: 702886821

README.md

@@ -135,8 +135,8 @@ Please read our [contributing guidelines](./CONTRIBUTING.md) and send us PRs!
 [Autograd]: https://github.com/hips/autograd "Autograd on GitHub"
 [DeepMind JAX Ecosystem]: https://deepmind.com/blog/article/using-jax-to-accelerate-our-research "DeepMind JAX Ecosystem"
 [DeepMind JAX Ecosystem citation]: https://github.com/deepmind/jax/blob/main/deepmind2020jax.txt "Citation"
-[JAX]: https://github.com/google/jax "JAX on GitHub"
-[JAX installation instructions]: https://github.com/google/jax#installation "JAX installation"
+[JAX]: https://github.com/jax-ml/jax "JAX on GitHub"
+[JAX installation instructions]: https://github.com/jax-ml/jax#installation "JAX installation"
 [jit]: https://jax.readthedocs.io/en/latest/jax.html#jax.jit "jax.jit documentation"
 [NumPy]: https://numpy.org/ "NumPy"
 [pmap]: https://jax.readthedocs.io/en/latest/jax.html#jax.pmap "jax.pmap documentation"

dm_pix/_src/augment_test.py

@@ -189,8 +189,8 @@ def test_flip(self, images_list):
         probability=(0., 1.))
 
   # Due to a bug in scipy we cannot test all available modes, refer to these
-  # issues for more information: https://github.com/google/jax/issues/11097,
-  # https://github.com/google/jax/issues/11097
+  # issues for more information: https://github.com/jax-ml/jax/issues/11097,
+  # https://github.com/jax-ml/jax/issues/11097
   @parameterized.named_parameters(
       ("in_range_nearest_0", _RAND_FLOATS_IN_RANGE, "nearest", 0),
       ("in_range_nearest_1", _RAND_FLOATS_IN_RANGE, "nearest", 1),