Merge branch 'main' into typ_v4

Author: crusaderky

array_api_compat/cupy/_aliases.py

@@ -118,6 +118,11 @@ def count_nonzero(
    return result
 
 
+# take_along_axis: axis defaults to -1 but in cupy (and numpy) axis is a required arg
+def take_along_axis(x: Array, indices: Array, /, *, axis: int = -1):
+    return cp.take_along_axis(x, indices, axis=axis)
+
+
 # These functions are completely new here. If the library already has them
 # (i.e., numpy 2.0), use the library version instead of our wrapper.
 if hasattr(cp, 'vecdot'):
@@ -139,6 +144,7 @@ def count_nonzero(
                               'acos', 'acosh', 'asin', 'asinh', 'atan',
                               'atan2', 'atanh', 'bitwise_left_shift',
                               'bitwise_invert', 'bitwise_right_shift',
-                              'bool', 'concat', 'count_nonzero', 'pow', 'sign']
+                              'bool', 'concat', 'count_nonzero', 'pow', 'sign',
+                              'take_along_axis']
 
 _all_ignore = ['cp', 'get_xp']

array_api_compat/numpy/_aliases.py

@@ -125,6 +125,11 @@ def count_nonzero(
     return result
 
 
+# take_along_axis: axis defaults to -1 but in numpy axis is a required arg
+def take_along_axis(x: Array, indices: Array, /, *, axis: int = -1):
+    return np.take_along_axis(x, indices, axis=axis)
+
+
 # These functions are completely new here. If the library already has them
 # (i.e., numpy 2.0), use the library version instead of our wrapper.
 if hasattr(np, "vecdot"):
@@ -160,6 +165,7 @@ def count_nonzero(
     "concat",
     "count_nonzero",
     "pow",
+    "take_along_axis"
 ]
 __all__ += _aliases.__all__
 _all_ignore = ["np", "get_xp"]

array_api_compat/torch/_aliases.py

@@ -3,7 +3,7 @@
 from collections.abc import Sequence
 from functools import reduce as _reduce, wraps as _wraps
 from builtins import all as _builtin_all, any as _builtin_any
-from typing import Any
+from typing import Any, List, Optional, Sequence, Tuple, Union, Literal
 
 import torch
 
@@ -547,8 +547,12 @@ def count_nonzero(
 ) -> Array:
     result = torch.count_nonzero(x, dim=axis)
     if keepdims:
-        if axis is not None:
+        if isinstance(axis, int):
             return result.unsqueeze(axis)
+        elif isinstance(axis, tuple):
+            n_axis = [x.ndim + ax if ax < 0 else ax for ax in axis]
+            sh = [1 if i in n_axis else x.shape[i] for i in range(x.ndim)]
+            return torch.reshape(result, sh)
         return _axis_none_keepdims(result, x.ndim, keepdims)
     else:
         return result
@@ -820,6 +824,12 @@ def sign(x: Array, /) -> Array:
         return out
 
 
+def meshgrid(*arrays: Array, indexing: Literal['xy', 'ij'] = 'xy') -> List[Array]:
+    # enforce the default of 'xy'
+    # TODO: is the return type a list or a tuple
+    return list(torch.meshgrid(*arrays, indexing='xy'))
+
+
 __all__ = ['__array_namespace_info__', 'asarray', 'result_type', 'can_cast',
            'permute_dims', 'bitwise_invert', 'newaxis', 'conj', 'add',
            'atan2', 'bitwise_and', 'bitwise_left_shift', 'bitwise_or',
@@ -836,6 +846,6 @@ def sign(x: Array, /) -> Array:
            'UniqueAllResult', 'UniqueCountsResult', 'UniqueInverseResult',
            'unique_all', 'unique_counts', 'unique_inverse', 'unique_values',
            'matmul', 'matrix_transpose', 'vecdot', 'tensordot', 'isdtype',
-           'take', 'take_along_axis', 'sign', 'finfo', 'iinfo', 'repeat']
+           'take', 'take_along_axis', 'sign', 'finfo', 'iinfo', 'repeat', 'meshgrid']
 
 _all_ignore = ['torch', 'get_xp']

cupy-xfails.txt

@@ -34,6 +34,11 @@ array_api_tests/test_operators_and_elementwise_functions.py::test_subtract[__sub
 array_api_tests/test_operators_and_elementwise_functions.py::test_add[__add__(x, s)]
 # floating point inaccuracy
 array_api_tests/test_operators_and_elementwise_functions.py::test_remainder[remainder(x1, x2)]
+# incomplete NEP50 support in CuPy 13.x (fixed in 14.0.0a1)
+array_api_tests/test_operators_and_elementwise_functions.py::test_binary_with_scalars_real[pow]
+array_api_tests/test_operators_and_elementwise_functions.py::test_binary_with_scalars_real[logaddexp]
+array_api_tests/test_operators_and_elementwise_functions.py::test_binary_with_scalars_real[nextafter]
+array_api_tests/test_operators_and_elementwise_functions.py::test_binary_with_scalars_real[hypot]
 
 # cupy (arg)min/max wrong with infinities
 # https://github.com/cupy/cupy/issues/7424
@@ -182,7 +187,6 @@ array_api_tests/test_signatures.py::test_func_signature[from_dlpack]
 array_api_tests/test_signatures.py::test_array_method_signature[__dlpack__]
 
 # 2024.12 support
-array_api_tests/test_signatures.py::test_func_signature[count_nonzero]
 array_api_tests/test_signatures.py::test_func_signature[bitwise_and]
 array_api_tests/test_signatures.py::test_func_signature[bitwise_left_shift]
 array_api_tests/test_signatures.py::test_func_signature[bitwise_or]

dask-xfails.txt

@@ -24,6 +24,9 @@ array_api_tests/test_creation_functions.py::test_linspace
 # Shape mismatch
 array_api_tests/test_indexing_functions.py::test_take
 
+# missing `take_along_axis`, https://github.com/dask/dask/issues/3663
+array_api_tests/test_indexing_functions.py::test_take_along_axis
+
 # Array methods and attributes not already on da.Array cannot be wrapped
 array_api_tests/test_has_names.py::test_has_names[array_method-__array_namespace__]
 array_api_tests/test_has_names.py::test_has_names[array_method-to_device]

numpy-1-22-xfails.txt

@@ -123,6 +123,20 @@ array_api_tests/test_operators_and_elementwise_functions.py::test_subtract[subtr
 array_api_tests/test_searching_functions.py::test_where
 array_api_tests/test_special_cases.py::test_iop[__iadd__(x1_i is -0 and x2_i is -0) -> -0]
 
+array_api_tests/test_operators_and_elementwise_functions.py::test_binary_with_scalars_real[add]
+array_api_tests/test_operators_and_elementwise_functions.py::test_binary_with_scalars_real[divide]
+array_api_tests/test_operators_and_elementwise_functions.py::test_binary_with_scalars_real[hypot]
+array_api_tests/test_operators_and_elementwise_functions.py::test_binary_with_scalars_real[subtract]
+array_api_tests/test_operators_and_elementwise_functions.py::test_binary_with_scalars_real[logaddexp]
+array_api_tests/test_operators_and_elementwise_functions.py::test_binary_with_scalars_real[nextafter]
+array_api_tests/test_operators_and_elementwise_functions.py::test_binary_with_scalars_real[multiply]
+array_api_tests/test_operators_and_elementwise_functions.py::test_binary_with_scalars_real[maximum]
+array_api_tests/test_operators_and_elementwise_functions.py::test_binary_with_scalars_real[copysign]
+array_api_tests/test_operators_and_elementwise_functions.py::test_binary_with_scalars_real[pow]
+array_api_tests/test_operators_and_elementwise_functions.py::test_binary_with_scalars_real[pow]
+
+array_api_tests/test_operators_and_elementwise_functions.py::test_where_with_scalars
+
 # 2023.12 support
 array_api_tests/test_has_names.py::test_has_names[creation-from_dlpack]
 array_api_tests/test_signatures.py::test_func_signature[from_dlpack]

numpy-1-26-xfails.txt

@@ -50,6 +50,8 @@ array_api_tests/test_signatures.py::test_func_signature[bitwise_right_shift]
 array_api_tests/test_signatures.py::test_func_signature[bitwise_xor]
 array_api_tests/test_data_type_functions.py::TestResultType::test_with_scalars
 
+array_api_tests/test_operators_and_elementwise_functions.py::test_where_with_scalars
+
 # Stubs have a comment: (**note**: libraries may return ``NaN`` to match Python behavior.); Apparently, NumPy does just that
 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]

tests/test_torch.py

@@ -100,3 +100,20 @@ def test_gh_273(self, default_dt, dtype_a, dtype_b):
             assert dtype_1 == dtype_2
         finally:
             torch.set_default_dtype(prev_default)
+
+
+def test_meshgrid():
+    """Verify that array_api_compat.torch.meshgrid defaults to indexing='xy'."""
+
+    x, y = xp.asarray([1, 2]), xp.asarray([4])
+
+    X, Y = xp.meshgrid(x, y)
+
+    # output of torch.meshgrid(x, y, indexing='xy') -- indexing='ij' is different
+    X_xy, Y_xy = xp.asarray([[1, 2]]), xp.asarray([[4, 4]])
+
+    assert X.shape == X_xy.shape
+    assert xp.all(X == X_xy)
+
+    assert Y.shape == Y_xy.shape
+    assert xp.all(Y == Y_xy)