Repository: jax-ml/jax
Pull request #27678 (2025-04-02 20:10:25)

jnp.concat and friends: support __jax_array__


Patch:
diff --git a/jax/_src/numpy/lax_numpy.py b/jax/_src/numpy/lax_numpy.py
index ae32703e7..43b892324 100644
--- a/jax/_src/numpy/lax_numpy.py
+++ b/jax/_src/numpy/lax_numpy.py
@@ -4466,7 +4466,7 @@ def stack(arrays: np.ndarray | Array | Sequence[ArrayLike],
     axis = _canonicalize_axis(axis, arrays.ndim)
     return concatenate(expand_dims(arrays, axis + 1), axis=axis, dtype=dtype)
   else:
-    util.check_arraylike("stack", *arrays)
+    arrays = util.ensure_arraylike_tuple("stack", arrays)
     shape0 = np.shape(arrays[0])
     axis = _canonicalize_axis(axis, len(shape0) + 1)
     new_arrays = []
@@ -4555,7 +4555,7 @@ def tile(A: ArrayLike, reps: DimSize | Sequence[DimSize]) -> Array:
            [1, 2],
            [3, 4]], dtype=int32)
   """
-  util.check_arraylike("tile", A)
+  A = util.ensure_arraylike("tile", A)
   try:
     iter(reps)  # type: ignore[arg-type]
   except TypeError:
@@ -4628,7 +4628,7 @@ def concatenate(arrays: np.ndarray | Array | Sequence[ArrayLike],
   """
   if isinstance(arrays, (np.ndarray, Array)):
     return _concatenate_array(arrays, axis, dtype=dtype)
-  util.check_arraylike("concatenate", *arrays)
+  arrays = util.ensure_arraylike_tuple("concatenate", arrays)
   if not len(arrays):
     raise ValueError("Need at least one array to concatenate.")
   if axis is None:
@@ -4870,6 +4870,7 @@ def dstack(tup: np.ndarray | Array | Sequence[ArrayLike],
   else:
     # TODO(jakevdp): Non-array input deprecated 2023-09-22; change to error.
     util.check_arraylike("dstack", *tup, emit_warning=True)
+    tup = util.ensure_arraylike_tuple("dstack", tup)
     arrs = [atleast_3d(m) for m in tup]
   return concatenate(arrs, axis=2, dtype=dtype)
 
diff --git a/jax/_src/numpy/util.py b/jax/_src/numpy/util.py
index e0e20d443..9d56267c4 100644
--- a/jax/_src/numpy/util.py
+++ b/jax/_src/numpy/util.py
@@ -159,7 +159,7 @@ def ensure_arraylike(fun_name: str, /, *args: Any) -> Array | tuple[Array, ...]:
   return tuple(_arraylike_asarray(arg) for arg in args)  # pytype: disable=bad-return-type
 
 
-def ensure_arraylike_tuple(fun_name: str, tup: tuple[Any, ...]) -> tuple[Array, ...]:
+def ensure_arraylike_tuple(fun_name: str, tup: Sequence[Any]) -> tuple[Array, ...]:
   """Check that argument elements are arraylike and convert to a tuple of arrays.
 
   This is useful because ensure_arraylike with a single argument returns a single array.


Test patch:
diff --git a/tests/array_extensibility_test.py b/tests/array_extensibility_test.py
index 55089720f..730001abe 100644
--- a/tests/array_extensibility_test.py
+++ b/tests/array_extensibility_test.py
@@ -267,10 +267,10 @@ NUMPY_APIS = [
   NumPyAPI.sig(jnp.ceil, Float[5]),
   # NumPyAPI.sig(jnp.choose, [int, float], [(3,), (10,)]),
   NumPyAPI.sig(jnp.clip, Float[5]),
-  # NumPyAPI.sig(jnp.column_stack, [float], [(3, 10)]),
+  NumPyAPI.sig(jnp.column_stack, [Float[5], Float[5], Float[5]]),
   NumPyAPI.sig(jnp.compress, Float[10], Bool[10]),
-  # NumPyAPI.sig(jnp.concat, [Float[5], Float[5]]),
-  # NumPyAPI.sig(jnp.concatenate, [Float[5], Float[5]]),
+  NumPyAPI.sig(jnp.concat, [Float[5], Float[5]]),
+  NumPyAPI.sig(jnp.concatenate, [Float[5], Float[5]]),
   NumPyAPI.sig(jnp.conj, Float[5]),
   NumPyAPI.sig(jnp.conjugate, Float[5]),
   NumPyAPI.sig(jnp.convolve, Float[7], Float[3]),
@@ -300,7 +300,7 @@ NUMPY_APIS = [
   NumPyAPI.sig(jnp.divmod, Float[5], Float[5]),
   NumPyAPI.sig(jnp.dot, Float[5], Float[5]),
   NumPyAPI.sig(jnp.dsplit, Float[3, 5, 6], indices_or_sections=2),
-  # NumPyAPI.sig(jnp.dstack, Float[3, 5]),
+  NumPyAPI.sig(jnp.dstack, [Float[3, 5, 1], Float[3, 5, 3]]),
   NumPyAPI.sig(jnp.ediff1d, Float[5]),
   NumPyAPI.sig(jnp.empty_like, Float[5]),
   NumPyAPI.sig(jnp.equal, Float[5], Float[5]),
@@ -469,7 +469,7 @@ NUMPY_APIS = [
   NumPyAPI.sig(jnp.sqrt, Float[5]),
   NumPyAPI.sig(jnp.square, Float[5]),
   NumPyAPI.sig(jnp.squeeze, Float[5]),
-  # NumPyAPI.sig(jnp.stack, [Float[2, 3], Float[2, 3]], axis=1),
+  NumPyAPI.sig(jnp.stack, [Float[2, 3], Float[2, 3]], axis=1),
   NumPyAPI.sig(jnp.std, Float[5]),
   NumPyAPI.sig(jnp.subtract, Float[5], Float[5]),
   NumPyAPI.sig(jnp.sum, Float[5]),
@@ -479,7 +479,7 @@ NUMPY_APIS = [
   NumPyAPI.sig(jnp.tan, Float[5]),
   NumPyAPI.sig(jnp.tanh, Float[5]),
   NumPyAPI.sig(jnp.tensordot, Float[2, 3, 4], Float[3, 4, 5]),
-  # NumPyAPI.sig(jnp.tile, Float[5], reps=(2,)),
+  NumPyAPI.sig(jnp.tile, Float[5], reps=(2,)),
   NumPyAPI.sig(jnp.trace, Float[5, 5]),
   NumPyAPI.sig(jnp.transpose, Float[5, 6]),
   NumPyAPI.sig(jnp.trapezoid, Float[5]),