Merge branch 'main' into clip

Author: crusaderky

.github/workflows/array-api-tests-dask.yml

@@ -10,4 +10,9 @@ jobs:
       package-version: '>= 2024.9.0'
       module-name: dask.array
       extra-requires: numpy
-      pytest-extra-args: --disable-deadline --max-examples=5
+      # Dask is substantially slower then other libraries on unit tests.
+      # Reduce the number of examples to speed up CI, even though this means that this
+      # workflow is barely more than a smoke test, and one should expect extreme 
+      # flakiness. Before changes to dask-xfails.txt or dask-skips.txt, please run
+      # the full test suite with at least 200 examples.
+      pytest-extra-args: --max-examples=5

.github/workflows/array-api-tests.yml

@@ -33,7 +33,7 @@ on:
         description: "Multiline string of environment variables to set for the test run."
 
 env:
-  PYTEST_ARGS: "--max-examples 200 -v -rxXfE --ci ${{ inputs.pytest-extra-args }} --hypothesis-disable-deadline"
+  PYTEST_ARGS: "--max-examples 200 -v -rxXfE --ci ${{ inputs.pytest-extra-args }} --hypothesis-disable-deadline --durations 10"
 
 jobs:
   tests:

.github/workflows/docs-deploy.yml

@@ -13,7 +13,7 @@ jobs:
     steps:
       - uses: actions/checkout@v4
       - name: Download Artifact
-        uses: dawidd6/action-download-artifact@v8
+        uses: dawidd6/action-download-artifact@v9
         with:
           workflow: docs-build.yml
           name: docs-build

array_api_compat/cupy/_aliases.py

@@ -125,6 +125,20 @@ def astype(
     return out.copy() if copy and out is x else out
 
 
+# cupy.count_nonzero does not have keepdims
+def count_nonzero(
+    x: ndarray,
+    axis=None,
+    keepdims=False
+) -> ndarray:
+   result = cp.count_nonzero(x, axis)
+   if keepdims:
+       if axis is None:
+            return cp.reshape(result, [1]*x.ndim)
+       return cp.expand_dims(result, axis)
+   return result
+
+
 # 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'):
@@ -146,6 +160,6 @@ def astype(
                               'acos', 'acosh', 'asin', 'asinh', 'atan',
                               'atan2', 'atanh', 'bitwise_left_shift',
                               'bitwise_invert', 'bitwise_right_shift',
-                              'bool', 'concat', 'pow', 'sign']
+                              'bool', 'concat', 'count_nonzero', 'pow', 'sign']
 
 _all_ignore = ['cp', 'get_xp']

array_api_compat/dask/array/_aliases.py

@@ -335,6 +335,21 @@ def argsort(
     return restore(x)
 
 
+# dask.array.count_nonzero does not have keepdims
+def count_nonzero(
+    x: Array,
+    axis=None,
+    keepdims=False
+) -> Array:
+   result = da.count_nonzero(x, axis)
+   if keepdims:
+       if axis is None:
+            return da.reshape(result, [1]*x.ndim)
+       return da.expand_dims(result, axis)
+   return result
+
+
+
 __all__ = _aliases.__all__ + [
                     '__array_namespace_info__', 'asarray', 'astype', 'acos',
                     'acosh', 'asin', 'asinh', 'atan', 'atan2',
@@ -343,6 +358,6 @@ def argsort(
                     'result_type', 'bool', 'float32', 'float64', 'int8', 'int16', 'int32', 'int64',
                     'uint8', 'uint16', 'uint32', 'uint64',
                     'complex64', 'complex128', 'iinfo', 'finfo',
-                    'can_cast', 'result_type']
+                    'can_cast', 'count_nonzero', 'result_type']
 
 _all_ignore = ["Callable", "array_namespace", "get_xp", "da", "np"]

array_api_compat/numpy/_aliases.py

@@ -127,6 +127,19 @@ def astype(
     return x.astype(dtype=dtype, copy=copy)
 
 
+# count_nonzero returns a python int for axis=None and keepdims=False
+# https://github.com/numpy/numpy/issues/17562
+def count_nonzero(
+    x : ndarray,
+    axis=None,
+    keepdims=False
+) -> ndarray:
+    result = np.count_nonzero(x, axis=axis, keepdims=keepdims)
+    if axis is None and not keepdims:
+        return np.asarray(result)
+    return result
+
+
 # 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'):
@@ -148,6 +161,6 @@ def astype(
                               'acos', 'acosh', 'asin', 'asinh', 'atan',
                               'atan2', 'atanh', 'bitwise_left_shift',
                               'bitwise_invert', 'bitwise_right_shift',
-                              'bool', 'concat', 'pow']
+                              'bool', 'concat', 'count_nonzero', 'pow']
 
 _all_ignore = ['np', 'get_xp']

array_api_compat/torch/_aliases.py

@@ -3,13 +3,7 @@
 from functools import wraps as _wraps
 from builtins import all as _builtin_all, any as _builtin_any
 
-from ..common._aliases import (matrix_transpose as _aliases_matrix_transpose,
-                               vecdot as _aliases_vecdot,
-                               clip as _aliases_clip,
-                               unstack as _aliases_unstack,
-                               cumulative_sum as _aliases_cumulative_sum,
-                               cumulative_prod as _aliases_cumulative_prod,
-                               )
+from ..common import _aliases
 from .._internal import get_xp
 
 from ._info import __array_namespace_info__
@@ -215,10 +209,10 @@ def min(x: array, /, *, axis: Optional[Union[int, Tuple[int, ...]]] = None, keep
         return torch.clone(x)
     return torch.amin(x, axis, keepdims=keepdims)
 
-clip = get_xp(torch)(_aliases_clip)
-unstack = get_xp(torch)(_aliases_unstack)
-cumulative_sum = get_xp(torch)(_aliases_cumulative_sum)
-cumulative_prod = get_xp(torch)(_aliases_cumulative_prod)
+clip = get_xp(torch)(_aliases.clip)
+unstack = get_xp(torch)(_aliases.unstack)
+cumulative_sum = get_xp(torch)(_aliases.cumulative_sum)
+cumulative_prod = get_xp(torch)(_aliases.cumulative_prod)
 
 # torch.sort also returns a tuple
 # https://github.com/pytorch/pytorch/issues/70921
@@ -527,15 +521,22 @@ def diff(
     return torch.diff(x, dim=axis, n=n, prepend=prepend, append=append)
 
 
-# torch uses `dim` instead of `axis`
+# torch uses `dim` instead of `axis`, does not have keepdims
 def count_nonzero(
     x: array,
     /,
     *,
     axis: Optional[Union[int, Tuple[int, ...]]] = None,
     keepdims: bool = False,
 ) -> array:
-    return torch.count_nonzero(x, dim=axis, keepdims=keepdims)
+    result = torch.count_nonzero(x, dim=axis)
+    if keepdims:
+        if axis is not None:
+            return result.unsqueeze(axis)
+        return _axis_none_keepdims(result, x.ndim, keepdims)
+    else:
+        return result
+
 
 
 def where(condition: array, x1: array, x2: array, /) -> array:
@@ -710,8 +711,8 @@ def matmul(x1: array, x2: array, /, **kwargs) -> array:
     x1, x2 = _fix_promotion(x1, x2, only_scalar=False)
     return torch.matmul(x1, x2, **kwargs)
 
-matrix_transpose = get_xp(torch)(_aliases_matrix_transpose)
-_vecdot = get_xp(torch)(_aliases_vecdot)
+matrix_transpose = get_xp(torch)(_aliases.matrix_transpose)
+_vecdot = get_xp(torch)(_aliases.vecdot)
 
 def vecdot(x1: array, x2: array, /, *, axis: int = -1) -> array:
     x1, x2 = _fix_promotion(x1, x2, only_scalar=False)

cupy-xfails.txt

@@ -191,5 +191,4 @@ array_api_tests/test_signatures.py::test_func_signature[bitwise_or]
 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_special_cases.py::test_binary[nextafter(x1_i is +0 and x2_i is -0) -> -0]
-array_api_tests/test_special_cases.py::test_nan_propagation[cumulative_prod]
 

dask-skips.txt

@@ -1,5 +1,9 @@
-# slow and not implemented in dask
-array_api_tests/test_linalg.py::test_matrix_power
+# NOTE: dask tests run on a very small number of examples in CI due to
+# slowness. This causes very high flakiness in the tests.
+# Before changing this file, please run with at least 200 examples.
 
-# hangs on 2024.12
+# Passes, but extremely slow
+array_api_tests/test_linalg.py::test_outer
+
+# Hangs
 array_api_tests/test_creation_functions.py::test_eye

dask-xfails.txt

@@ -1,3 +1,11 @@
+# NOTE: dask tests run on a very small number of examples in CI due to
+# slowness. This causes very high flakiness in the tests.
+# Before changing this file, please run with at least 200 examples.
+
+# Broken edge case with shape 0
+# https://github.com/dask/dask/issues/11800
+array_api_tests/test_array_object.py::test_setitem
+
 # Various indexing errors
 array_api_tests/test_array_object.py::test_getitem_masking
 
@@ -11,6 +19,9 @@ array_api_tests/test_data_type_functions.py::test_finfo[float32]
 # (I think the test is not forcing the op to be computed?)
 array_api_tests/test_creation_functions.py::test_linspace
 
+# Shape mismatch
+array_api_tests/test_indexing_functions.py::test_take
+
 # 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]
@@ -31,7 +42,7 @@ array_api_tests/test_set_functions.py::test_unique_values
 # fails for ndim > 2
 array_api_tests/test_linalg.py::test_svdvals
 
-# dtype mismatch got uint64, but should be uint8, NPY_PROMOTION_STATE=weak doesn't help :(
+# dtype mismatch got uint64, but should be uint8; NPY_PROMOTION_STATE=weak doesn't help
 array_api_tests/test_linalg.py::test_tensordot
 
 # AssertionError: out.dtype=uint64, but should be uint8 [tensordot(uint8, uint8)]
@@ -89,29 +100,29 @@ array_api_tests/meta/test_hypothesis_helpers.py::test_symmetric_matrices
 # https://github.com/dask/dask/issues/11706
 array_api_tests/test_creation_functions.py::test_arange
 
+# da.searchsorted with a sorter argument is not supported
+array_api_tests/test_searching_functions.py::test_searchsorted
+
 # 2023.12 support
 array_api_tests/test_manipulation_functions.py::test_repeat
 
 # 2024.12 support
-array_api_tests/test_array_object.py::test_setitem
 array_api_tests/test_array_object.py::test_getitem_arrays_and_ints_1[1]
 array_api_tests/test_array_object.py::test_getitem_arrays_and_ints_1[None]
 array_api_tests/test_array_object.py::test_getitem_arrays_and_ints_2[1]
 array_api_tests/test_array_object.py::test_getitem_arrays_and_ints_2[None]
 array_api_tests/test_has_names.py::test_has_names[indexing-take_along_axis]
+array_api_tests/test_signatures.py::test_func_signature[count_nonzero]
 array_api_tests/test_signatures.py::test_func_signature[take_along_axis]
 
 array_api_tests/test_linalg.py::test_cholesky
 array_api_tests/test_linalg.py::test_linalg_matmul
+array_api_tests/test_linalg.py::test_matmul
 array_api_tests/test_linalg.py::test_matrix_norm
 array_api_tests/test_linalg.py::test_qr
-array_api_tests/test_manipulation_functions.py::test_concat
 array_api_tests/test_manipulation_functions.py::test_roll
-array_api_tests/test_operators_and_elementwise_functions.py::test_add[add(x1, x2)]
-array_api_tests/test_operators_and_elementwise_functions.py::test_bitwise_left_shift[bitwise_left_shift(x1, x2)]
-array_api_tests/test_operators_and_elementwise_functions.py::test_bitwise_right_shift[bitwise_right_shift(x1, x2)]
-array_api_tests/test_operators_and_elementwise_functions.py::test_greater[__gt__(x1, x2)]
-array_api_tests/test_signatures.py::test_func_signature[count_nonzero]
+
+# Stubs have a comment: (**note**: libraries may return ``NaN`` to match Python behavior.)
 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[floor_divide(x1_i is -infinity and isfinite(x2_i) and x2_i > 0) -> -infinity]
@@ -130,6 +141,3 @@ array_api_tests/test_special_cases.py::test_iop[__ifloordiv__(x1_i is -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__(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_nan_propagation[cumulative_prod]
-
-

docs/changelog.md

@@ -1,6 +1,27 @@
 # Changelog
 
-## 1.11.0 (2025-XX-XX)
+## 1.11.1 (2025-03-04)
+
+This is a bugfix release with no new features compared to version 1.11.
+
+### Major Changes
+
+- fix `count_nonzero` wrappers: work around the lack of the `keepdims` argument in
+  several array libraries (torch, dask, cupy); work around numpy returning python
+  ints in for some input combinations.
+
+### Minor Changes
+
+- runnings self-tests does not require all array libraries. Missing libraries are
+  skipped.
+
+The following users contributed to this release:
+
+Evgeni Burovski
+Guido Imperiale
+
+
+## 1.11.0 (2025-02-27)
 
 ### Major Changes
 

numpy-1-21-xfails.txt

@@ -187,13 +187,13 @@ array_api_tests/test_signatures.py::test_array_method_signature[__dlpack__]
 array_api_tests/test_manipulation_functions.py::test_repeat
 
 # 2024.12 support
-array_api_tests/test_special_cases.py::test_nan_propagation[cumulative_prod]
-
 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]
 array_api_tests/test_signatures.py::test_func_signature[bitwise_right_shift]
 array_api_tests/test_signatures.py::test_func_signature[bitwise_xor]
+
+# 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]
 array_api_tests/test_special_cases.py::test_binary[floor_divide(x1_i is -infinity and isfinite(x2_i) and x2_i > 0) -> -infinity]

numpy-1-26-xfails.txt

@@ -41,13 +41,13 @@ array_api_tests/test_signatures.py::test_array_method_signature[__dlpack__]
 array_api_tests/test_manipulation_functions.py::test_repeat
 
 # 2024.12 support
-array_api_tests/test_special_cases.py::test_nan_propagation[cumulative_prod]
-
 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]
 array_api_tests/test_signatures.py::test_func_signature[bitwise_right_shift]
 array_api_tests/test_signatures.py::test_func_signature[bitwise_xor]
+
+# 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]
 array_api_tests/test_special_cases.py::test_binary[floor_divide(x1_i is -infinity and isfinite(x2_i) and x2_i > 0) -> -infinity]

numpy-dev-xfails.txt

@@ -11,13 +11,13 @@ array_api_tests/test_signatures.py::test_extension_func_signature[linalg.vecdot]
 array_api_tests/test_manipulation_functions.py::test_repeat
 
 # 2024.12 support
-array_api_tests/test_special_cases.py::test_nan_propagation[cumulative_prod]
-
 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]
 array_api_tests/test_signatures.py::test_func_signature[bitwise_right_shift]
 array_api_tests/test_signatures.py::test_func_signature[bitwise_xor]
+
+# 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]
 array_api_tests/test_special_cases.py::test_binary[floor_divide(x1_i is -infinity and isfinite(x2_i) and x2_i > 0) -> -infinity]

numpy-xfails.txt

@@ -13,8 +13,6 @@ array_api_tests/test_signatures.py::test_array_method_signature[__dlpack__]
 array_api_tests/test_manipulation_functions.py::test_repeat
 
 # 2024.12 support
-array_api_tests/test_special_cases.py::test_nan_propagation[cumulative_prod]
-
 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]

tests/_helpers.py

@@ -8,8 +8,7 @@
 ]
 
 def import_(library, wrapper=False):
-    if library in ('cupy', 'ndonnx'):
-        pytest.importorskip(library)
+    pytest.importorskip(library)
     if wrapper:
         if 'jax' in library:
             # JAX v0.4.32 implements the array API directly in jax.numpy

torch-xfails.txt

@@ -136,4 +136,3 @@ array_api_tests/test_signatures.py::test_array_method_signature[__lshift__]
 array_api_tests/test_signatures.py::test_array_method_signature[__or__]
 array_api_tests/test_signatures.py::test_array_method_signature[__rshift__]
 array_api_tests/test_signatures.py::test_array_method_signature[__xor__]
-array_api_tests/test_special_cases.py::test_nan_propagation[cumulative_prod]