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: 733951719
Change-Id: Iee4a7bf4bcc9fc3108ca9d30b6b98f5df20c79c7

Author: jakeharmon8

mjx/README.md

@@ -9,7 +9,7 @@
 
 This package is a re-implementation of the
 [MuJoCo physics engine](https://github.com/google-deepmind/mujoco) in
-[JAX](https://github.com/google/jax). This library is developed and maintained
+[JAX](https://github.com/jax-ml/jax). This library is developed and maintained
 by Google DeepMind, and is kept up-to-date with the latest developments in
 MuJoCo itself.
 

mjx/mujoco/mjx/_src/math.py

@@ -51,7 +51,7 @@ def norm(
   """Calculates a linalg.norm(x) that's safe for gradients at x=0.
 
   Avoids a poorly defined gradient for jnp.linal.norm(0) see
-  https://github.com/google/jax/issues/3058 for details
+  https://github.com/jax-ml/jax/issues/3058 for details
   Args:
     x: A jnp.array
     axis: The axis along which to compute the norm