Use clear dtype annotations to avoid dtype promotion issues

The default dtype changes depending on the x64 mode so it's better
to be explicit in code internal to JAX.

PiperOrigin-RevId: 801797558

Author: apaszke

jax/_src/lax_reference.py

@@ -285,7 +285,8 @@ def ragged_dot(
 
   out = np.zeros((m, n), dtype=lhs.dtype)
   result_iota = np.expand_dims(np.arange(out.shape[0]), list(range(1, out.ndim)))
-  start = 0
+  result_iota = result_iota.astype(group_sizes.dtype)
+  start = np.asarray(0, dtype=group_sizes.dtype)
   for i, size in enumerate(group_sizes):
     out += np.where(
         np.logical_and(start <= result_iota, result_iota < (start + size)),

jax/_src/numpy/lax_numpy.py

@@ -7631,7 +7631,7 @@ def delete(
     obj = asarray(obj).ravel()
     obj = clip(where(obj < 0, obj + a.shape[axis], obj), 0, a.shape[axis])
     obj = sort(obj)
-    obj -= arange(len(obj))  # type: ignore[arg-type,operator]
+    obj -= arange(len(obj), dtype=obj.dtype)  # type: ignore
     i = arange(a.shape[axis] - obj.size)
     i += (i[None, :] >= obj[:, None]).sum(0)
     return a[(slice(None),) * axis + (i,)]

tests/lax_numpy_test.py

@@ -1655,7 +1655,10 @@ def testIntegerPower(self, ptype):
   def testIntegerPowerOverflow(self, x, y):
     # Regression test for https://github.com/jax-ml/jax/issues/5987
     args_maker = lambda: [x, y]
-    self._CheckAgainstNumpy(np.power, jnp.power, args_maker)
+    check_dtypes = platform.system() != 'Windows'
+    self._CheckAgainstNumpy(
+        np.power, jnp.power, args_maker, check_dtypes=check_dtypes
+    )
     self._CompileAndCheck(jnp.power, args_maker)
 
   @jtu.sample_product(
@@ -3852,7 +3855,11 @@ def testArrayFromList(self):
 
     # out of bounds leads to an OverflowError
     val = jnp.iinfo(jnp.int64).min - 1
-    with self.assertRaisesRegex(OverflowError, "Python int too large.*"):
+    if platform.system() == 'Windows':
+      expected_regex = 'int too big to convert'
+    else:
+      expected_regex = 'Python int too large.*'
+    with self.assertRaisesRegex(OverflowError, expected_regex):
       jnp.array([0, val])
 
   def testArrayNone(self):

tests/state_test.py

@@ -506,9 +506,9 @@ def f(x_ref, *idxs):
       return op(x_ref, indexer)
 
     rng = self.rng()
-    a = rng.randn(*bat_ref_aval.shape)
+    a = rng.randn(*bat_ref_aval.shape).astype(floatx)
     his = [d for d, b in zip(ref_aval.shape, indexed_dims) if b]
-    idxs = [rng.randint(low=0, high=hi, size=i.shape)
+    idxs = [rng.randint(low=0, high=hi, size=i.shape, dtype=intx)
             for i, hi in zip(bat_idx_avals, his)]
 
     # discharge-of-vmap