Repository: jax-ml/jax
Pull request #27716 (2025-04-03 22:18:21)

jax.numpy: support __jax_array__ in remaining APIs


Patch:
diff --git a/.github/workflows/k8s.yaml b/.github/workflows/k8s.yaml
index 1042388fe..a96ce1ead 100644
--- a/.github/workflows/k8s.yaml
+++ b/.github/workflows/k8s.yaml
@@ -38,7 +38,7 @@ jobs:
           path: jax
 
       - name: Start Minikube cluster
-        uses: medyagh/setup-minikube@cea33675329b799adccc9526aa5daccc26cd5052 # ratchet:medyagh/setup-minikube@v0.0.19
+        uses: medyagh/setup-minikube@d8c0eb871f6f455542491d86a574477bd3894533 # ratchet:medyagh/setup-minikube@v0.0.18
 
       - name: Install K8s Jobset
         run: |
diff --git a/jax/_src/dispatch.py b/jax/_src/dispatch.py
index baab6d519..d205f860b 100644
--- a/jax/_src/dispatch.py
+++ b/jax/_src/dispatch.py
@@ -44,7 +44,6 @@ from jax._src.interpreters import mlir
 from jax._src.interpreters import pxla
 from jax._src.interpreters import xla
 from jax._src.layout import DeviceLocalLayout, Layout
-from jax._src.lib import jaxlib_extension_version
 from jax._src.lib import xla_client as xc
 from jax._src.mesh import AbstractMesh, Mesh
 from jax._src.monitoring import record_event_duration_secs, record_event_time_span
@@ -496,7 +495,7 @@ def _device_put_sharding_impl(x, aval, device, copy):
         return _DeferredShardArg(x, x.sharding, aval, x.committed, copy)
     elif is_single_device_sharding(x.sharding):
       device = x.sharding._device_assignment[0] if device is None else device
-      if copy == CopySemantics.COPY and jaxlib_extension_version >= 327:
+      if copy == CopySemantics.COPY:
         return xc.batched_device_put(aval, SingleDeviceSharding(device), [x],
                                      [device], True, True)
       return pxla.batched_device_put(aval, SingleDeviceSharding(device), [x],
diff --git a/jax/_src/numpy/lax_numpy.py b/jax/_src/numpy/lax_numpy.py
index d42266170..503dca178 100644
--- a/jax/_src/numpy/lax_numpy.py
+++ b/jax/_src/numpy/lax_numpy.py
@@ -911,11 +911,11 @@ def histogram(a: ArrayLike, bins: ArrayLike = 10,
     Array(True, dtype=bool)
   """
   if weights is None:
-    util.check_arraylike("histogram", a, bins)
+    a, _ = util.ensure_arraylike("histogram", a, bins)
     a, = util.promote_dtypes_inexact(a)
     weights = ones_like(a)
   else:
-    util.check_arraylike("histogram", a, bins, weights)
+    a, _, weights = util.ensure_arraylike("histogram", a, bins, weights)
     if np.shape(a) != np.shape(weights):
       raise ValueError("weights should have the same shape as a.")
     a, weights = util.promote_dtypes_inexact(a, weights)
@@ -1005,7 +1005,7 @@ def histogram2d(x: ArrayLike, y: ArrayLike, bins: ArrayLike | list[ArrayLike] =
     >>> jnp.allclose(normed_sum, 1.0)
     Array(True, dtype=bool)
   """
-  util.check_arraylike("histogram2d", x, y)
+  x, y = util.ensure_arraylike("histogram2d", x, y)
   try:
     N = len(bins)  # type: ignore[arg-type]
   except TypeError:
@@ -1077,10 +1077,10 @@ def histogramdd(sample: ArrayLike, bins: ArrayLike | list[ArrayLike] = 10,
     Array(True, dtype=bool)
   """
   if weights is None:
-    util.check_arraylike("histogramdd", sample)
+    sample = util.ensure_arraylike("histogramdd", sample)
     sample, = util.promote_dtypes_inexact(sample)
   else:
-    util.check_arraylike("histogramdd", sample, weights)
+    sample, weights = util.ensure_arraylike("histogramdd", sample, weights)
     if np.shape(weights) != np.shape(sample)[:1]:
       raise ValueError("should have one weight for each sample.")
     sample, weights = util.promote_dtypes_inexact(sample, weights)
@@ -2424,7 +2424,7 @@ def expand_dims(a: ArrayLike, axis: int | Sequence[int]) -> Array:
              [2],
              [3]]]], dtype=int32)
   """
-  util.check_arraylike("expand_dims", a)
+  a = util.ensure_arraylike("expand_dims", a)
   axis = _ensure_index_tuple(axis)
   return lax.expand_dims(a, axis)
 
@@ -4371,7 +4371,7 @@ def pad(array: ArrayLike, pad_width: PadValueLike[int | Array | np.ndarray],
     Array([-10, -10,   2,   3,   4,  10,  10], dtype=int32)
   """
 
-  util.check_arraylike("pad", array)
+  array = util.ensure_arraylike("pad", array)
   pad_width = _broadcast_to_pairs(pad_width, np.ndim(array), "pad_width")
   if pad_width and not all(core.is_dim(p[0]) and core.is_dim(p[1])
                            for p in pad_width):
@@ -6988,8 +6988,10 @@ def repeat(a: ArrayLike, repeats: ArrayLike, axis: int | None = None, *,
     Array([[1, 1, 2, 2, 2, 2, 2],
            [3, 3, 4, 4, 4, 4, 4]], dtype=int32)
   """
-  arr = util.ensure_arraylike("repeat", a)
-  core.is_dim(repeats) or util.check_arraylike("repeat", repeats)
+  if core.is_dim(repeats):
+    arr = util.ensure_arraylike("repeat", a)
+  else:
+    arr, repeats = util.ensure_arraylike("repeat", a, repeats)
 
   if axis is None:
     arr = arr.ravel()
@@ -7828,7 +7830,7 @@ def diag_indices_from(arr: ArrayLike) -> tuple[Array, ...]:
     Array([0, 1], dtype=int32),
     Array([0, 1], dtype=int32))
   """
-  util.check_arraylike("diag_indices_from", arr)
+  arr = util.ensure_arraylike("diag_indices_from", arr)
   nd = np.ndim(arr)
   if not np.ndim(arr) >= 2:
     raise ValueError("input array must be at least 2-d")
@@ -8244,6 +8246,9 @@ def delete(
   # Case 3: obj is an array
   # NB: pass both arrays to check for appropriate error message.
   util.check_arraylike("delete", a, obj)
+  # Can't use ensure_arraylike here because obj may be static.
+  if hasattr(obj, "__jax_array__"):
+    obj = obj.__jax_array__()
 
   # Case 3a: unique integer indices; delete in a JIT-compatible way
   if issubdtype(_dtype(obj), np.integer) and assume_unique_indices:
diff --git a/jax/experimental/mosaic/gpu/utils.py b/jax/experimental/mosaic/gpu/utils.py
index 47401440f..28534cf40 100644
--- a/jax/experimental/mosaic/gpu/utils.py
+++ b/jax/experimental/mosaic/gpu/utils.py
@@ -135,7 +135,7 @@ def _debug_scalar_ty_format(arg):
     return "%llu", arg
   if ir.F32Type.isinstance(arg.type):
     return "%f", arg
-  if ir.BF16Type.isinstance(arg.type) or ir.F16Type.isinstance(arg.type):
+  if ir.F16Type.isinstance(arg.type):
     arg = arith.extf(ir.F32Type.get(), arg)
     return "%f", arg
   raise NotImplementedError(f"Can't print the type {arg.type}")


Test patch:
diff --git a/tests/array_extensibility_test.py b/tests/array_extensibility_test.py
index 14fcc18ca..7c2681a21 100644
--- a/tests/array_extensibility_test.py
+++ b/tests/array_extensibility_test.py
@@ -81,6 +81,7 @@ class ShapeDtype:
 
 Bool = ShapeDtype(bool)
 Int = ShapeDtype(int)
+UInt = ShapeDtype('uint32')
 Uint8 = ShapeDtype('uint8')
 Float = ShapeDtype(float)
 Complex = ShapeDtype(complex)
@@ -280,18 +281,18 @@ NUMPY_APIS = [
   NumPyAPI.sig(jnp.correlate, Float[7], Float[3]),
   NumPyAPI.sig(jnp.cos, Float[5]),
   NumPyAPI.sig(jnp.cosh, Float[5]),
-  # NumPyAPI.sig(np.count_nonzero, [float], [(10,)]),
-  # NumPyAPI.sig(np.cov, [float], [(10,)]),
-  # NumPyAPI.sig(np.cross, [float, float], [(3,), (3,)]),
+  NumPyAPI.sig(jnp.count_nonzero, Float[10]),
+  NumPyAPI.sig(jnp.cov, Float[10]),
+  NumPyAPI.sig(jnp.cross, Float[3], Float[3]),
   NumPyAPI.sig(jnp.cumprod, Float[5]),
   NumPyAPI.sig(jnp.cumsum, Float[5]),
   NumPyAPI.sig(jnp.cumulative_prod, Float[5]),
   NumPyAPI.sig(jnp.cumulative_sum, Float[5]),
   NumPyAPI.sig(jnp.deg2rad, Float[5]),
   NumPyAPI.sig(jnp.degrees, Float[5]),
-  # NumPyAPI.sig(jnp.delete, Float[5], Int[()]),
+  NumPyAPI.sig(jnp.delete, Float[5], Int[()]),
   NumPyAPI.sig(jnp.diag, Float[5]),
-  # NumPyAPI.sig(jnp.diag_indices_from, Float[5, 5]),
+  NumPyAPI.sig(jnp.diag_indices_from, Float[5, 5]),
   NumPyAPI.sig(jnp.diagflat, Float[5]),
   NumPyAPI.sig(jnp.diagonal, Float[5, 5]),
   NumPyAPI.sig(jnp.diff, Float[5]),
@@ -306,7 +307,7 @@ NUMPY_APIS = [
   NumPyAPI.sig(jnp.equal, Float[5], Float[5]),
   NumPyAPI.sig(jnp.exp, Float[5]),
   NumPyAPI.sig(jnp.exp2, Float[5]),
-  # NumPyAPI.sig(jnp.expand_dims, Float[5], axis=0),
+  NumPyAPI.sig(jnp.expand_dims, Float[5], axis=0),
   NumPyAPI.sig(jnp.expm1, Float[5]),
   NumPyAPI.sig(jnp.extract, Bool[5], Float[5]),
   NumPyAPI.sig(jnp.fabs, Float[5]),
@@ -332,11 +333,11 @@ NUMPY_APIS = [
   NumPyAPI.sig(jnp.greater, Float[5], Float[5]),
   NumPyAPI.sig(jnp.greater_equal, Float[5], Float[5]),
   NumPyAPI.sig(jnp.heaviside, Float[5], Float[5]),
-  # NumPyAPI.sig(jnp.histogram, Float[5]),
+  NumPyAPI.sig(jnp.histogram, Float[5]),
   NumPyAPI.sig(jnp.histogram2d, Float[5], Float[5]),
   NumPyAPI.sig(jnp.histogram_bin_edges, Float[5]),
-  # NumPyAPI.sig(jnp.histogramdd, Float[5, 3]),
-  # NumPyAPI.sig(jnp.hsplit, Float[3, 5], Int[1]),
+  NumPyAPI.sig(jnp.histogramdd, Float[5, 3]),
+  NumPyAPI.sig(jnp.hsplit, Float[3, 6], indices_or_sections=2),
   NumPyAPI.sig(jnp.hstack, (Float[5], Float[5])),
   NumPyAPI.sig(jnp.hypot, Float[5], Float[5]),
   NumPyAPI.sig(jnp.i0, Float[5]),
@@ -411,7 +412,7 @@ NUMPY_APIS = [
   NumPyAPI.sig(jnp.ones_like, Float[5]),
   NumPyAPI.sig(jnp.outer, Float[5], Float[5]),
   NumPyAPI.sig(jnp.packbits, Int[5]),
-  # NumPyAPI.sig(jnp.pad, Float[5], pad_width=2),
+  NumPyAPI.sig(jnp.pad, Float[5], pad_width=2),
   NumPyAPI.sig(jnp.partition, Float[5], kth=3),
   NumPyAPI.sig(jnp.percentile, Float[5], q=75),
   NumPyAPI.sig(jnp.permute_dims, Float[3, 5], axes=(1, 0)),
@@ -437,11 +438,11 @@ NUMPY_APIS = [
   NumPyAPI.sig(jnp.rad2deg, Float[5]),
   NumPyAPI.sig(jnp.radians, Float[5]),
   NumPyAPI.sig(jnp.ravel, Float[5]),
-  # NumPyAPI.sig(jnp.ravel_multi_index, Int[2, 5], dims=(2, 3)),
+  NumPyAPI.sig(jnp.ravel_multi_index, [Uint8[5], Uint8[5]], dims=(8, 9)),
   NumPyAPI.sig(jnp.real, Complex[5]),
   NumPyAPI.sig(jnp.reciprocal, Float[5]),
   NumPyAPI.sig(jnp.remainder, Float[5], Float[5]),
-  # NumPyAPI.sig(jnp.repeat, Float[5], Int[5]),
+  NumPyAPI.sig(jnp.repeat, Float[5], Int[5]),
   NumPyAPI.sig(jnp.reshape, Float[6], shape=(2, 3)),
   NumPyAPI.sig(jnp.resize, Float[6], new_shape=(2, 3)),
   NumPyAPI.sig(jnp.right_shift, Int[5], Int[5]),
diff --git a/tests/pjit_test.py b/tests/pjit_test.py
index 2db75be18..2570c6090 100644
--- a/tests/pjit_test.py
+++ b/tests/pjit_test.py
@@ -1401,7 +1401,7 @@ class PJitTest(jtu.BufferDonationTestCase):
     self.assertIn("stablehlo.constant dense<-0.000000e+00>", ir)
 
   def test_device_put_copy_donate(self):
-    if jaxlib_extension_version < 327:
+    if jaxlib_extension_version < 323:
       raise unittest.SkipTest("Copy not supported in device put.")
     x = np.arange(1000)
     y = jax.device_put(x, device=jax.devices()[0], may_alias=False, donate=False)