Repository: jax-ml/jax
Pull request #27394 (2025-03-24 23:01:53)

Support __jax_array__ in jnp.full_like & co


Patch:
diff --git a/jax/_src/numpy/array_creation.py b/jax/_src/numpy/array_creation.py
index 67418e732..a0495986f 100644
--- a/jax/_src/numpy/array_creation.py
+++ b/jax/_src/numpy/array_creation.py
@@ -244,6 +244,8 @@ def zeros_like(a: ArrayLike | DuckTypedArray,
            [0, 0, 0]], dtype=int32)
   """
   if not (hasattr(a, 'dtype') and hasattr(a, 'shape')):  # support duck typing
+    if hasattr(a, '__jax_array__'):
+      a = a.__jax_array__()
     util.check_arraylike("zeros_like", a)
   dtypes.check_user_dtype_supported(dtype, "zeros_like")
   if shape is not None:
@@ -287,6 +289,8 @@ def ones_like(a: ArrayLike | DuckTypedArray,
            [1, 1, 1]], dtype=int32)
   """
   if not (hasattr(a, 'dtype') and hasattr(a, 'shape')):  # support duck typing
+    if hasattr(a, '__jax_array__'):
+      a = a.__jax_array__()
     util.check_arraylike("ones_like", a)
   dtypes.check_user_dtype_supported(dtype, "ones_like")
   if shape is not None:
@@ -332,9 +336,13 @@ def empty_like(prototype: ArrayLike | DuckTypedArray,
            [0, 0, 0]], dtype=int32)
   """
   if not (hasattr(prototype, 'dtype') and hasattr(prototype, 'shape')):  # support duck typing
-    util.check_arraylike("empty_like", prototype)
-  dtypes.check_user_dtype_supported(dtype, "empty_like")
-  return zeros_like(prototype, dtype=dtype, shape=shape, device=device)
+    if hasattr(prototype, '__jax_array__'):
+      prototype = prototype.__jax_array__()
+    util.check_arraylike("ones_like", prototype)
+  dtypes.check_user_dtype_supported(dtype, "ones_like")
+  if shape is not None:
+    shape = canonicalize_shape(shape)
+  return lax.full_like(prototype, 0, dtype, shape, sharding=util.normalize_device_to_sharding(device))
 
 
 @export
@@ -382,6 +390,8 @@ def full_like(a: ArrayLike | DuckTypedArray,
     util.check_arraylike("full_like", 0, fill_value)
   else:
     util.check_arraylike("full_like", a, fill_value)
+    if hasattr(a, '__jax_array__'):
+      a = a.__jax_array__()
   dtypes.check_user_dtype_supported(dtype, "full_like")
   if shape is not None:
     shape = canonicalize_shape(shape)
diff --git a/jax/_src/pallas/mosaic_gpu/primitives.py b/jax/_src/pallas/mosaic_gpu/primitives.py
index 8eafa0ac8..fe28766cf 100644
--- a/jax/_src/pallas/mosaic_gpu/primitives.py
+++ b/jax/_src/pallas/mosaic_gpu/primitives.py
@@ -489,7 +489,6 @@ def _barrier_arrive_abstract_eval(barrier, *args, **params):
 
 
 @lowering.register_lowering_rule(barrier_arrive_p, mgpu.ThreadSemantics.Lane)
-@lowering.register_lowering_rule(barrier_arrive_p, mgpu.ThreadSemantics.Warpgroup)
 def _barrier_arrive_lowering(
     ctx: lowering.LoweringRuleContext,
     barrier,


Test patch:
diff --git a/tests/array_extensibility_test.py b/tests/array_extensibility_test.py
index 45c83f796..3e84f6668 100644
--- a/tests/array_extensibility_test.py
+++ b/tests/array_extensibility_test.py
@@ -12,10 +12,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import functools
+from typing import Any, Callable, NamedTuple
 
 from absl.testing import absltest
 from absl.testing import parameterized
-from typing import Any, Callable, NamedTuple
+import numpy as np
 
 import jax
 import jax.numpy as jnp
@@ -38,6 +40,15 @@ class JaxArrayWrapper:
     return jnp.asarray(self.x)
 
 
+class DuckTypedArrayWithErroringJaxArray:
+  """Duck-typed array that provides a __jax_array__ method which fails."""
+  shape = (2, 3)
+  dtype = np.dtype('float32')
+
+  def __jax_array__(self):
+    raise ValueError("jax array was called.")
+
+
 class NumPyAPI(NamedTuple):
   fun: Callable[..., Any]
   args: list[jax.ShapeDtypeStruct]
@@ -287,7 +298,7 @@ NUMPY_APIS = [
   NumPyAPI.sig(jnp.dsplit, Float[3, 5, 6], indices_or_sections=2),
   # NumPyAPI.sig(jnp.dstack, Float[3, 5]),
   NumPyAPI.sig(jnp.ediff1d, Float[5]),
-  # NumPyAPI.sig(jnp.empty_like, Float[5]),
+  NumPyAPI.sig(jnp.empty_like, Float[5]),
   NumPyAPI.sig(jnp.equal, Float[5], Float[5]),
   NumPyAPI.sig(jnp.exp, Float[5]),
   NumPyAPI.sig(jnp.exp2, Float[5]),
@@ -312,7 +323,7 @@ NUMPY_APIS = [
   NumPyAPI.sig(jnp.fmin, Float[5], Float[5]),
   NumPyAPI.sig(jnp.fmod, Float[5], Float[5]),
   NumPyAPI.sig(jnp.frexp, Float[5]),
-  # NumPyAPI.sig(jnp.full_like, Float[5], Float[()]),
+  NumPyAPI.sig(jnp.full_like, Float[5], Float[()]),
   NumPyAPI.sig(jnp.gcd, Int[5], Int[5]),
   NumPyAPI.sig(jnp.greater, Float[5], Float[5]),
   NumPyAPI.sig(jnp.greater_equal, Float[5], Float[5]),
@@ -393,7 +404,7 @@ NUMPY_APIS = [
   NumPyAPI.sig(jnp.nextafter, Float[5], Float[5]),
   NumPyAPI.sig(jnp.nonzero, Float[5]),
   NumPyAPI.sig(jnp.not_equal, Float[5], Float[5]),
-  # NumPyAPI.sig(jnp.ones_like, Float[5]),
+  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),
@@ -493,7 +504,7 @@ NUMPY_APIS = [
   NumPyAPI.sig(jnp.vsplit, Float[6], indices_or_sections=2),
   NumPyAPI.sig(jnp.vstack, [Float[5], Float[2, 5]]),
   NumPyAPI.sig(jnp.where, Bool[5], Float[5], Float[5]),
-  # NumPyAPI.sig(jnp.zeros_like, Float[5]),
+  NumPyAPI.sig(jnp.zeros_like, Float[5]),
 ]
 
 
@@ -511,6 +522,26 @@ class JaxArrayTests(jtu.JaxTestCase):
 
     self.assertAllClose(wrapped, expected, atol=0, rtol=0)
 
+  @parameterized.named_parameters(
+    {'testcase_name': func.__name__, 'func': func}
+    for func in [jnp.zeros_like, jnp.ones_like, jnp.empty_like, jnp.full_like]
+  )
+  def test_array_creation_from_duck_typed_array(self, func):
+    # Ensure that jnp.*_like prefers shape/dtype over __jax_array__ when
+    # both methods are available.
+    if func is jnp.full_like:
+      func = functools.partial(func, fill_value=2.0)
+    obj = DuckTypedArrayWithErroringJaxArray()
+
+    # The test relies on this failing
+    with self.assertRaises(ValueError):
+      jnp.asarray(obj)
+
+    result = func(obj)
+    self.assertIsInstance(result, jax.Array)
+    self.assertEqual(result.shape, obj.shape)
+    self.assertEqual(result.dtype, obj.dtype)
+
 
 if __name__ == "__main__":
   absltest.main(testLoader=jtu.JaxTestLoader())