Fixing CI test failures

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

@@ -24,7 +24,7 @@ on:
 
 
 env:
-  PYTEST_ARGS: "--max-examples 200 -v -rxXfE --ci ${{ inputs.pytest-extra-args }}"
+  PYTEST_ARGS: "--max-examples 200 -v -rxXfE --ci ${{ inputs.pytest-extra-args }} --hypothesis-disable-deadline"
 
 jobs:
   tests:

array_api_compat/common/_aliases.py

@@ -386,6 +386,12 @@ def sort(
         res = xp.flip(res, axis=axis)
     return res
 
+# nonzero should error for zero-dimensional arrays
+def nonzero(x: ndarray, /, xp, **kwargs) -> Tuple[ndarray, ...]:
+    if x.ndim == 0:
+        raise ValueError("nonzero() does not support zero-dimensional arrays")
+    return xp.nonzero(x, **kwargs)
+
 # sum() and prod() should always upcast when dtype=None
 def sum(
     x: ndarray,
@@ -526,5 +532,5 @@ def isdtype(
            'UniqueAllResult', 'UniqueCountsResult', 'UniqueInverseResult',
            'unique_all', 'unique_counts', 'unique_inverse', 'unique_values',
            'astype', 'std', 'var', 'permute_dims', 'reshape', 'argsort',
-           'sort', 'sum', 'prod', 'ceil', 'floor', 'trunc', 'matmul',
-           'matrix_transpose', 'tensordot', 'vecdot', 'isdtype']
+           'sort', 'nonzero', 'sum', 'prod', 'ceil', 'floor', 'trunc',
+           'matmul', 'matrix_transpose', 'tensordot', 'vecdot', 'isdtype']

array_api_compat/cupy/_aliases.py

@@ -52,6 +52,7 @@
 reshape = get_xp(cp)(_aliases.reshape)
 argsort = get_xp(cp)(_aliases.argsort)
 sort = get_xp(cp)(_aliases.sort)
+nonzero = get_xp(cp)(_aliases.nonzero)
 sum = get_xp(cp)(_aliases.sum)
 prod = get_xp(cp)(_aliases.prod)
 ceil = get_xp(cp)(_aliases.ceil)

array_api_compat/numpy/_aliases.py

@@ -52,6 +52,7 @@
 reshape = get_xp(np)(_aliases.reshape)
 argsort = get_xp(np)(_aliases.argsort)
 sort = get_xp(np)(_aliases.sort)
+nonzero = get_xp(np)(_aliases.nonzero)
 sum = get_xp(np)(_aliases.sum)
 prod = get_xp(np)(_aliases.prod)
 ceil = get_xp(np)(_aliases.ceil)

array_api_compat/torch/_aliases.py

@@ -475,6 +475,8 @@ def roll(x: array, /, shift: Union[int, Tuple[int, ...]], *, axis: Optional[Unio
     return torch.roll(x, shift, axis, **kwargs)
 
 def nonzero(x: array, /, **kwargs) -> Tuple[array, ...]:
+    if x.ndim == 0:
+        raise ValueError("nonzero() does not support zero-dimensional arrays")
     return torch.nonzero(x, as_tuple=True, **kwargs)
 
 def where(condition: array, x1: array, x2: array, /) -> array:

array_api_compat/torch/linalg.py

@@ -4,6 +4,8 @@
 if TYPE_CHECKING:
     import torch
     array = torch.Tensor
+    from torch import dtype as Dtype
+    from typing import Optional
 
 from torch.linalg import *
 
@@ -12,9 +14,9 @@
 from torch import linalg as torch_linalg
 linalg_all = [i for i in dir(torch_linalg) if not i.startswith('_')]
 
-# These are implemented in torch but aren't in the linalg namespace
-from torch import outer, trace
-from ._aliases import _fix_promotion, matrix_transpose, tensordot
+# outer is implemented in torch but aren't in the linalg namespace
+from torch import outer
+from ._aliases import _fix_promotion, matrix_transpose, tensordot, sum
 
 # Note: torch.linalg.cross does not default to axis=-1 (it defaults to the
 # first axis with size 3), see https://github.com/pytorch/pytorch/issues/58743
@@ -49,6 +51,11 @@ def solve(x1: array, x2: array, /, **kwargs) -> array:
     x1, x2 = _fix_promotion(x1, x2, only_scalar=False)
     return torch.linalg.solve(x1, x2, **kwargs)
 
+# torch.trace doesn't support the offset argument and doesn't support stacking
+def trace(x: array, /, *, offset: int = 0, dtype: Optional[Dtype] = None) -> array:
+    # Use our wrapped sum to make sure it does upcasting correctly
+    return sum(torch.diagonal(x, offset=offset, dim1=-2, dim2=-1), axis=-1, dtype=dtype)
+
 __all__ = linalg_all + ['outer', 'trace', 'matrix_transpose', 'tensordot',
                         'vecdot', 'solve']
 

cupy-xfails.txt

@@ -175,3 +175,16 @@ array_api_tests/test_special_cases.py::test_iop[__imod__(isfinite(x1_i) and x1_i
 array_api_tests/test_special_cases.py::test_iop[__imod__(isfinite(x1_i) and x1_i > 0 and x2_i is -infinity) -> x2_i]
 array_api_tests/test_special_cases.py::test_iop[__imod__(isfinite(x1_i) and x1_i < 0 and x2_i is +infinity) -> x2_i]
 array_api_tests/test_special_cases.py::test_iop[__imod__(isfinite(x1_i) and x1_i < 0 and x2_i is -infinity) -> x1_i]
+
+# fft functions are not yet supported
+# (https://github.com/data-apis/array-api-compat/issues/67)
+array_api_tests/test_fft.py::test_fft
+array_api_tests/test_fft.py::test_ifft
+array_api_tests/test_fft.py::test_fftn
+array_api_tests/test_fft.py::test_ifftn
+array_api_tests/test_fft.py::test_rfft
+array_api_tests/test_fft.py::test_irfft
+array_api_tests/test_fft.py::test_rfftn
+array_api_tests/test_fft.py::test_irfftn
+array_api_tests/test_fft.py::test_hfft
+array_api_tests/test_fft.py::test_ihfft

numpy-1-21-xfails.txt

@@ -236,3 +236,16 @@ array_api_tests/test_special_cases.py::test_binary[remainder(x1_i is +0 and x2_i
 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[remainder(x1_i is -0 and x2_i > 0) -> +0]
 array_api_tests/test_special_cases.py::test_iop[__iadd__(x1_i is -0 and x2_i is -0) -> -0]
+
+# fft functions are not yet supported
+# (https://github.com/data-apis/array-api-compat/issues/67)
+array_api_tests/test_fft.py::test_fft
+array_api_tests/test_fft.py::test_ifft
+array_api_tests/test_fft.py::test_fftn
+array_api_tests/test_fft.py::test_ifftn
+array_api_tests/test_fft.py::test_rfft
+array_api_tests/test_fft.py::test_irfft
+array_api_tests/test_fft.py::test_rfftn
+array_api_tests/test_fft.py::test_irfftn
+array_api_tests/test_fft.py::test_hfft
+array_api_tests/test_fft.py::test_ihfft

numpy-xfails.txt

@@ -54,3 +54,16 @@ array_api_tests/meta/test_hypothesis_helpers.py::test_symmetric_matrices
 # The test suite is incorrectly checking sums that have loss of significance
 # (https://github.com/data-apis/array-api-tests/issues/168)
 array_api_tests/test_statistical_functions.py::test_sum
+
+# fft functions are not yet supported
+# (https://github.com/data-apis/array-api-compat/issues/67)
+array_api_tests/test_fft.py::test_fft
+array_api_tests/test_fft.py::test_ifft
+array_api_tests/test_fft.py::test_fftn
+array_api_tests/test_fft.py::test_ifftn
+array_api_tests/test_fft.py::test_rfft
+array_api_tests/test_fft.py::test_irfft
+array_api_tests/test_fft.py::test_rfftn
+array_api_tests/test_fft.py::test_irfftn
+array_api_tests/test_fft.py::test_hfft
+array_api_tests/test_fft.py::test_ihfft