Merge pull request #33688 from jakevdp:fix-dep-dtype

PiperOrigin-RevId: 839817597

Author: Google-ML-Automation

jax/_src/lax/lax.py

@@ -3485,7 +3485,7 @@ def _delta(dtype: DTypeLike, shape: Shape, axes: Sequence[int]) -> Array:
 def _tri(dtype: DTypeLike, shape: Shape, offset: DimSize) -> Array:
   """Like numpy.tri, create a 2D array with ones below a diagonal."""
   offset = asarray(core.dimension_as_value(offset))
-  if not dtypes.issubdtype(offset, np.integer):
+  if not dtypes.issubdtype(offset.dtype, np.integer):
     raise TypeError(f"offset must be an integer, got {offset!r}")
   shape_dtype = lax_utils.int_dtype_for_shape(shape, signed=True)
   if (

jax/_src/numpy/reductions.py

@@ -202,7 +202,7 @@ def _cast_to_numeric(operand: Array) -> Array:
   return promote_dtypes_numeric(operand)[0]
 
 def _require_integer(arr: Array) -> Array:
-  if not dtypes.isdtype(arr, ("bool", "integral")):
+  if not dtypes.isdtype(arr.dtype, ("bool", "integral")):
     raise ValueError(f"integer argument required; got dtype={arr.dtype}")
   return arr
 

tests/api_test.py

@@ -4174,7 +4174,7 @@ def __jax_array__(self):
 
     x = jnp.array(1)
     a = AlexArray(x)
-    for f in [jnp.isscalar, jnp.size, jnp.shape, jnp.dtype]:
+    for f in [jnp.isscalar, jnp.size, jnp.shape, jnp.result_type]:
       self.assertEqual(f(x), f(a))
 
     x = AlexArray(jnp.array(1))