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: 702886818

ml_dtypes/tests/custom_float_test.py

@@ -241,7 +241,7 @@ def testModuleName(self, float_type):
 
   @ignore_warning(category=RuntimeWarning, message="invalid value encountered")
   def testPickleable(self, float_type):
-    # https://github.com/google/jax/discussions/8505
+    # https://github.com/jax-ml/jax/discussions/8505
     x = np.arange(10, dtype=float_type)
     serialized = pickle.dumps(x)
     x_out = pickle.loads(serialized)
@@ -694,7 +694,7 @@ def testHash(self, float_type):
     self.assertEqual(h, hash(np.dtype(float_type.__name__)))
 
   def testDeepCopyDoesNotAlterHash(self, float_type):
-    # For context, see https://github.com/google/jax/issues/4651. If the hash
+    # For context, see https://github.com/jax-ml/jax/issues/4651. If the hash
     # value of the type descriptor is not initialized correctly, a deep copy
     # can change the type hash.
     dtype = np.dtype(float_type)

ml_dtypes/tests/intn_test.py

@@ -59,7 +59,7 @@ def testModuleName(self, scalar_type):
 
   @parameterized.product(scalar_type=INTN_TYPES)
   def testPickleable(self, scalar_type):
-    # https://github.com/google/jax/discussions/8505
+    # https://github.com/jax-ml/jax/discussions/8505
     x = np.arange(10, dtype=scalar_type)
     serialized = pickle.dumps(x)
     x_out = pickle.loads(serialized)
@@ -264,7 +264,7 @@ def testHash(self, scalar_type):
 
   @parameterized.product(scalar_type=INTN_TYPES)
   def testDeepCopyDoesNotAlterHash(self, scalar_type):
-    # For context, see https://github.com/google/jax/issues/4651. If the hash
+    # For context, see https://github.com/jax-ml/jax/issues/4651. If the hash
     # value of the type descriptor is not initialized correctly, a deep copy
     # can change the type hash.
     dtype = np.dtype(scalar_type)