Repository: jax-ml/jax
Pull request #27681 (2025-04-02 21:29:26)

jax.numpy: support __jax_array__ in several more functions


Patch:
diff --git a/.github/workflows/jax-array-api.yml b/.github/workflows/jax-array-api.yml
index 7df4228dd..c91ab6b8b 100644
--- a/.github/workflows/jax-array-api.yml
+++ b/.github/workflows/jax-array-api.yml
@@ -28,7 +28,7 @@ jobs:
       with:
         repository: data-apis/array-api-tests
         # TODO(jakevdp) update this to a stable release/tag when available.
-        ref: 'c48410f96fc58e02eea844e6b7f6cc01680f77ce'  # Latest commit as of 2025-04-02
+        ref: '0b89c5268e4e4a352223a487b8f63dbd1023872d'  # Latest commit as of 2025-03-04
         submodules: 'true'
         path: 'array-api-tests'
     - name: Set up Python ${{ matrix.python-version }}
diff --git a/jax/_src/numpy/array_api_metadata.py b/jax/_src/numpy/array_api_metadata.py
index d634a2856..5267e5121 100644
--- a/jax/_src/numpy/array_api_metadata.py
+++ b/jax/_src/numpy/array_api_metadata.py
@@ -29,7 +29,7 @@ from jax._src import dtypes as _dtypes
 from jax._src import xla_bridge as xb
 
 
-__array_api_version__ = '2024.12'
+__array_api_version__ = '2023.12'
 
 
 def __array_namespace__(self, *, api_version: None | str = None) -> ModuleType:
diff --git a/jax/_src/numpy/lax_numpy.py b/jax/_src/numpy/lax_numpy.py
index 43b892324..dba208327 100644
--- a/jax/_src/numpy/lax_numpy.py
+++ b/jax/_src/numpy/lax_numpy.py
@@ -2984,7 +2984,7 @@ def bincount(x: ArrayLike, weights: ArrayLike | None = None,
     >>> jnp.bincount(x, length=5)
     Array([2, 1, 0, 1, 0], dtype=int32)
   """
-  util.check_arraylike("bincount", x)
+  x = util.ensure_arraylike("bincount", x)
   if _dtype(x) == bool:
     x = lax.convert_element_type(x, 'int32')
   if not issubdtype(_dtype(x), np.integer):
@@ -5018,7 +5018,7 @@ def choose(a: ArrayLike, choices: Array | np.ndarray | Sequence[ArrayLike],
   """
   if out is not None:
     raise NotImplementedError("The 'out' argument to jnp.choose is not supported.")
-  util.check_arraylike('choose', a, *choices)
+  a, *choices = util.ensure_arraylike_tuple('choose', (a, *choices))
   if not issubdtype(_dtype(a), np.integer):
     raise ValueError("`a` array must be integer typed")
   N = len(choices)
@@ -8781,6 +8781,7 @@ def argwhere(
     >>> jnp.argwhere(0)
     Array([], shape=(0, 0), dtype=int32)
   """
+  a = util.ensure_arraylike("argwhere", a)
   result = transpose(vstack(nonzero(atleast_1d(a), size=size, fill_value=fill_value)))
   if np.ndim(a) == 0:
     return result[:0].reshape(result.shape[0], 0)
diff --git a/jax/_src/numpy/util.py b/jax/_src/numpy/util.py
index 9d56267c4..49605ffc3 100644
--- a/jax/_src/numpy/util.py
+++ b/jax/_src/numpy/util.py
@@ -259,7 +259,7 @@ def _broadcast_arrays(*args: ArrayLike) -> list[Array]:
 
 def _broadcast_to(arr: ArrayLike, shape: DimSize | Shape, sharding=None
                   ) -> Array:
-  check_arraylike("broadcast_to", arr)
+  arr = ensure_arraylike("broadcast_to", arr)
   arr = arr if isinstance(arr, Array) else lax.asarray(arr)
   if not isinstance(shape, tuple) and np.ndim(shape) == 0:
     shape = (shape,)


Test patch:
diff --git a/tests/array_api_skips.txt b/tests/array_api_skips.txt
index 7534cf6f8..2f8d4d1c6 100644
--- a/tests/array_api_skips.txt
+++ b/tests/array_api_skips.txt
@@ -10,24 +10,6 @@ array_api_tests/test_creation_functions.py::test_asarray_arrays
 
 # Returns wrong zero sign
 array_api_tests/test_special_cases.py::test_unary[sign((x_i is -0 or x_i == +0)) -> 0]
-array_api_tests/test_special_cases.py::test_iop[__imod__(x1_i is +0 and x2_i < 0) -> -0]
-array_api_tests/test_special_cases.py::test_iop[__imod__(x1_i is -0 and x2_i > 0) -> +0]
-array_api_tests/test_special_cases.py::test_binary[floor_divide(x1_i is -0 and x2_i > 0) -> -0]
-array_api_tests/test_special_cases.py::test_binary[floor_divide(x1_i is -0 and x2_i < 0) -> +0]
-array_api_tests/test_special_cases.py::test_binary[__floordiv__(x1_i is -0 and x2_i > 0) -> -0]
-array_api_tests/test_special_cases.py::test_binary[__floordiv__(isfinite(x1_i) and x1_i > 0 and x2_i is -infinity) -> -0]
-array_api_tests/test_special_cases.py::test_binary[__floordiv__(x1_i is -0 and x2_i < 0) -> +0]
-array_api_tests/test_special_cases.py::test_binary[__floordiv__(isfinite(x1_i) and x1_i < 0 and x2_i is +infinity) -> -0]
-array_api_tests/test_special_cases.py::test_binary[remainder(x1_i is -0 and x2_i > 0) -> +0]
-array_api_tests/test_special_cases.py::test_binary[__mod__(x1_i is -0 and x2_i > 0) -> +0]
-array_api_tests/test_special_cases.py::test_binary[floor_divide(isfinite(x1_i) and x1_i < 0 and x2_i is -infinity) -> +0]
-array_api_tests/test_special_cases.py::test_binary[__mod__(x1_i is +0 and x2_i < 0) -> -0]
-array_api_tests/test_special_cases.py::test_binary[remainder(x1_i is +0 and x2_i < 0) -> -0]
-array_api_tests/test_special_cases.py::test_binary[__floordiv__(isfinite(x1_i) and x1_i < 0 and x2_i is -infinity) -> +0]
-array_api_tests/test_special_cases.py::test_iop[__ifloordiv__(x1_i is -0 and x2_i > 0) -> -0]
-array_api_tests/test_special_cases.py::test_iop[__ifloordiv__(x1_i is -0 and x2_i < 0) -> +0]
-array_api_tests/test_special_cases.py::test_iop[__ifloordiv__(isfinite(x1_i) and x1_i < 0 and x2_i is -infinity) -> +0]
-
 
 # Returns int32 when int64 is expected
 array_api_tests/test_searching_functions.py::test_searchsorted
@@ -37,40 +19,3 @@ array_api_tests/test_operators_and_elementwise_functions.py::test_clip
 
 # JAX raises a ValueError rather than the expected IndexError for out-of-bound axis
 array_api_tests/test_manipulation_functions.py::test_expand_dims
-
-# Doesn't promote to uint64
-array_api_tests/test_statistical_functions.py::test_cumulative_prod
-
-# TODO(jakevdp): fix the following failures:
-
-# Returns NaN rather than inf
-array_api_tests/test_special_cases.py::test_binary[floor_divide(x1_i is +infinity and isfinite(x2_i) and x2_i < 0) -> -infinity]
-array_api_tests/test_special_cases.py::test_binary[floor_divide(x1_i is -infinity and isfinite(x2_i) and x2_i > 0) -> -infinity]
-array_api_tests/test_special_cases.py::test_binary[__floordiv__(x1_i < 0 and x2_i is +0) -> -infinity]
-array_api_tests/test_special_cases.py::test_binary[floor_divide(x1_i > 0 and x2_i is +0) -> +infinity]
-array_api_tests/test_special_cases.py::test_binary[floor_divide(x1_i is -infinity and isfinite(x2_i) and x2_i < 0) -> +infinity]
-array_api_tests/test_special_cases.py::test_binary[__floordiv__(x1_i < 0 and x2_i is -0) -> +infinity]
-array_api_tests/test_special_cases.py::test_binary[floor_divide(isfinite(x1_i) and x1_i > 0 and x2_i is -infinity) -> -0]
-array_api_tests/test_special_cases.py::test_binary[floor_divide(x1_i > 0 and x2_i is -0) -> -infinity]
-array_api_tests/test_special_cases.py::test_binary[__floordiv__(x1_i is +infinity and isfinite(x2_i) and x2_i > 0) -> +infinity]
-array_api_tests/test_special_cases.py::test_binary[floor_divide(x1_i < 0 and x2_i is +0) -> -infinity]
-array_api_tests/test_special_cases.py::test_binary[__floordiv__(x1_i is +infinity and isfinite(x2_i) and x2_i < 0) -> -infinity]
-array_api_tests/test_special_cases.py::test_binary[floor_divide(x1_i < 0 and x2_i is -0) -> +infinity]
-array_api_tests/test_special_cases.py::test_binary[__floordiv__(x1_i is -infinity and isfinite(x2_i) and x2_i > 0) -> -infinity]
-array_api_tests/test_special_cases.py::test_binary[floor_divide(x1_i is +infinity and isfinite(x2_i) and x2_i > 0) -> +infinity]
-array_api_tests/test_special_cases.py::test_binary[__floordiv__(x1_i is -infinity and isfinite(x2_i) and x2_i < 0) -> +infinity]
-array_api_tests/test_special_cases.py::test_binary[__floordiv__(x1_i > 0 and x2_i is +0) -> +infinity]
-array_api_tests/test_special_cases.py::test_binary[__floordiv__(x1_i > 0 and x2_i is -0) -> -infinity]
-array_api_tests/test_special_cases.py::test_iop[__ifloordiv__(x1_i > 0 and x2_i is +0) -> +infinity]
-array_api_tests/test_special_cases.py::test_iop[__ifloordiv__(x1_i is +infinity and isfinite(x2_i) and x2_i > 0) -> +infinity]
-array_api_tests/test_special_cases.py::test_iop[__ifloordiv__(x1_i > 0 and x2_i is -0) -> -infinity]
-array_api_tests/test_special_cases.py::test_iop[__ifloordiv__(x1_i is +infinity and isfinite(x2_i) and x2_i < 0) -> -infinity]
-array_api_tests/test_special_cases.py::test_iop[__ifloordiv__(x1_i is -infinity and isfinite(x2_i) and x2_i > 0) -> -infinity]
-array_api_tests/test_special_cases.py::test_iop[__ifloordiv__(x1_i < 0 and x2_i is -0) -> +infinity]
-array_api_tests/test_special_cases.py::test_iop[__ifloordiv__(x1_i is -infinity and isfinite(x2_i) and x2_i < 0) -> +infinity]
-array_api_tests/test_special_cases.py::test_iop[__ifloordiv__(x1_i < 0 and x2_i is +0) -> -infinity]
-
-# Returns -1.0 rather than 0.0
-array_api_tests/test_special_cases.py::test_binary[floor_divide(isfinite(x1_i) and x1_i < 0 and x2_i is +infinity) -> -0]
-array_api_tests/test_special_cases.py::test_iop[__ifloordiv__(isfinite(x1_i) and x1_i > 0 and x2_i is -infinity) -> -0]
-array_api_tests/test_special_cases.py::test_iop[__ifloordiv__(isfinite(x1_i) and x1_i < 0 and x2_i is +infinity) -> -0]
\ No newline at end of file
diff --git a/tests/array_extensibility_test.py b/tests/array_extensibility_test.py
index 730001abe..14fcc18ca 100644
--- a/tests/array_extensibility_test.py
+++ b/tests/array_extensibility_test.py
@@ -92,6 +92,8 @@ SKIPPED_APIS = [
  'apply_along_axis',
  'apply_over_axes',
  'arange',
+ 'array_str',
+ 'array_repr',
  'astype',
  'bartlett',
  'bfloat16',
@@ -101,6 +103,7 @@ SKIPPED_APIS = [
  'bool_',
  'broadcast_shapes',
  'c_',
+ 'can_cast',
  'cdouble',
  'character',
  'complex128',
@@ -233,14 +236,12 @@ NUMPY_APIS = [
   NumPyAPI.sig(jnp.argmin, Float[10]),
   NumPyAPI.sig(jnp.argpartition, Float[10], kth=5),
   NumPyAPI.sig(jnp.argsort, Float[10]),
-  # NumPyAPI.sig(jnp.argwhere, [float], [(10,)]),
+  NumPyAPI.sig(jnp.argwhere, Float[10]),
   NumPyAPI.sig(jnp.around, Float[5]),
   NumPyAPI.sig(jnp.array, Float[5]),
   NumPyAPI.sig(jnp.array_equal, Float[5], Float[5]),
   NumPyAPI.sig(jnp.array_equiv, Float[5], Float[5]),
-  # NumPyAPI.sig(jnp.array_repr, Float[5]),
   NumPyAPI.sig(jnp.array_split, Float[9], indices_or_sections=3),
-  # NumPyAPI.sig(jnp.array_str, Float[5]),
   NumPyAPI.sig(jnp.asarray, Float[5]),
   NumPyAPI.sig(jnp.asin, Float[5]),
   NumPyAPI.sig(jnp.asinh, Float[5]),
@@ -251,7 +252,7 @@ NUMPY_APIS = [
   NumPyAPI.sig(jnp.atleast_2d, Float[5]),
   NumPyAPI.sig(jnp.atleast_3d, Float[5]),
   NumPyAPI.sig(jnp.average, Float[10]),
-  # NumPyAPI.sig(jnp.bincount, int[10]),
+  NumPyAPI.sig(jnp.bincount, Int[10]),
   NumPyAPI.sig(jnp.bitwise_and, Int[5], Int[5]),
   NumPyAPI.sig(jnp.bitwise_count, Int[5]),
   NumPyAPI.sig(jnp.bitwise_invert, Int[5]),
@@ -261,11 +262,10 @@ NUMPY_APIS = [
   NumPyAPI.sig(jnp.bitwise_right_shift, Int[5], Int[5]),
   NumPyAPI.sig(jnp.bitwise_xor, Int[5], Int[5]),
   NumPyAPI.sig(jnp.broadcast_arrays, Float[5]),
-  # NumPyAPI.sig(jnp.broadcast_to, Float[()], shape=(10,)),
-  # NumPyAPI.sig(jnp.can_cast, Float[()], to='int32'),
+  NumPyAPI.sig(jnp.broadcast_to, Float[()], shape=(10,)),
   NumPyAPI.sig(jnp.cbrt, Float[5]),
   NumPyAPI.sig(jnp.ceil, Float[5]),
-  # NumPyAPI.sig(jnp.choose, [int, float], [(3,), (10,)]),
+  NumPyAPI.sig(jnp.choose, Int[3], [Float[3], Float[3], Float[3]], mode='clip'),
   NumPyAPI.sig(jnp.clip, Float[5]),
   NumPyAPI.sig(jnp.column_stack, [Float[5], Float[5], Float[5]]),
   NumPyAPI.sig(jnp.compress, Float[10], Bool[10]),