More improvements to test_linalg

array_api_tests/array_helpers.py

@@ -7,8 +7,12 @@
 # These are exported here so that they can be included in the special cases
 # tests from this file.
 from ._array_module import logical_not, subtract, floor, ceil, where
+from . import _array_module as xp
 from . import dtype_helpers as dh
 
+from ndindex import iter_indices
+
+import math
 
 __all__ = ['all', 'any', 'logical_and', 'logical_or', 'logical_not', 'less',
            'less_equal', 'greater', 'subtract', 'negative', 'floor', 'ceil',
@@ -146,6 +150,43 @@ def exactly_equal(x, y):
 
     return equal(x, y)
 
+def allclose(x, y, rel_tol=0.25, abs_tol=1, return_indices=False):
+    """
+    Return True all elements of x and y are within tolerance
+
+    If return_indices=True, returns (False, (i, j)) when the arrays are not
+    close, where i and j are the indices into x and y of corresponding
+    non-close elements.
+    """
+    for i, j in iter_indices(x.shape, y.shape):
+        i, j = i.raw, j.raw
+        a = x[i]
+        b = y[j]
+        if not (math.isfinite(a) and math.isfinite(b)):
+            # TODO: If a and b are both infinite, require the same type of infinity
+            continue
+        close = math.isclose(a, b, rel_tol=rel_tol, abs_tol=abs_tol)
+        if not close:
+            if return_indices:
+                return (False, (i, j))
+            return False
+    return True
+
+def assert_allclose(x, y, rel_tol=1, abs_tol=0.):
+    """
+    Test that x and y are approximately equal to each other.
+
+    Also asserts that x and y have the same shape and dtype.
+    """
+    assert x.shape == y.shape, f"The input arrays do not have the same shapes ({x.shape} != {y.shape})"
+
+    assert x.dtype == y.dtype, f"The input arrays do not have the same dtype ({x.dtype} != {y.dtype})"
+
+    c = allclose(x, y, rel_tol=rel_tol, abs_tol=abs_tol, return_indices=True)
+    if c is not True:
+        _, (i, j) = c
+        raise AssertionError(f"The input arrays are not close with {rel_tol = } and {abs_tol = } at indices {i = } and {j = }")
+
 def notequal(x, y):
     """
     Same as not_equal(x, y) except it gives False when both values are nan.
@@ -164,19 +205,21 @@ def notequal(x, y):
 
     return not_equal(x, y)
 
-def assert_exactly_equal(x, y):
+def assert_exactly_equal(x, y, msg_extra=None):
     """
     Test that the arrays x and y are exactly equal.
 
     If x and y do not have the same shape and dtype, they are not considered
     equal.
 
     """
-    assert x.shape == y.shape, f"The input arrays do not have the same shapes ({x.shape} != {y.shape})"
+    extra = '' if not msg_extra else f' ({msg_extra})'
 
-    assert x.dtype == y.dtype, f"The input arrays do not have the same dtype ({x.dtype} != {y.dtype})"
+    assert x.shape == y.shape, f"The input arrays do not have the same shapes ({x.shape} != {y.shape}){extra}"
+
+    assert x.dtype == y.dtype, f"The input arrays do not have the same dtype ({x.dtype} != {y.dtype}){extra}"
 
-    assert all(exactly_equal(x, y)), "The input arrays have different values"
+    assert all(exactly_equal(x, y)), f"The input arrays have different values ({x!r} != {y!r}){extra}"
 
 def assert_finite(x):
     """
@@ -306,3 +349,13 @@ def same_sign(x, y):
 def assert_same_sign(x, y):
     assert all(same_sign(x, y)), "The input arrays do not have the same sign"
 
+def _matrix_transpose(x):
+    if not isinstance(xp.matrix_transpose, xp._UndefinedStub):
+        return xp.matrix_transpose(x)
+    if hasattr(x, 'mT'):
+        return x.mT
+    if not isinstance(xp.permute_dims, xp._UndefinedStub):
+        perm = list(range(x.ndim))
+        perm[-1], perm[-2] = perm[-2], perm[-1]
+        return xp.permute_dims(x, axes=tuple(perm))
+    raise NotImplementedError("No way to compute matrix transpose")

array_api_tests/dtype_helpers.py

@@ -231,6 +231,57 @@ class MinMax(NamedTuple):
     {"complex64": xp.float32, "complex128": xp.float64}
 )
 
+def as_real_dtype(dtype):
+    """
+    Return the corresponding real dtype for a given floating-point dtype.
+    """
+    if dtype in real_float_dtypes:
+        return dtype
+    elif dtype_to_name[dtype] in complex_names:
+        return dtype_components[dtype]
+    else:
+        raise ValueError("as_real_dtype requires a floating-point dtype")
+
+def accumulation_result_dtype(x_dtype, dtype_kwarg):
+    """
+    Result dtype logic for sum(), prod(), and trace()
+
+    Note: may return None if a default uint cannot exist (e.g., for pytorch
+    which doesn't support uint32 or uint64). See https://github.com/data-apis/array-api-tests/issues/106
+
+    """
+    if dtype_kwarg is None:
+        if is_int_dtype(x_dtype):
+            if x_dtype in uint_dtypes:
+                default_dtype = default_uint
+            else:
+                default_dtype = default_int
+            if default_dtype is None:
+                _dtype = None
+            else:
+                m, M = dtype_ranges[x_dtype]
+                d_m, d_M = dtype_ranges[default_dtype]
+                if m < d_m or M > d_M:
+                    _dtype = x_dtype
+                else:
+                    _dtype = default_dtype
+        elif is_float_dtype(x_dtype, include_complex=False):
+            if dtype_nbits[x_dtype] > dtype_nbits[default_float]:
+                _dtype = x_dtype
+            else:
+                _dtype = default_float
+        elif api_version > "2021.12":
+            # Complex dtype
+            if dtype_nbits[x_dtype] > dtype_nbits[default_complex]:
+                _dtype = x_dtype
+            else:
+                _dtype = default_complex
+        else:
+            raise RuntimeError("Unexpected dtype. This indicates a bug in the test suite.")
+    else:
+        _dtype = dtype_kwarg
+
+    return _dtype
 
 if not hasattr(xp, "asarray"):
     default_int = xp.int32

array_api_tests/hypothesis_helpers.py

@@ -12,6 +12,7 @@
                                    sampled_from, shared, builds)
 
 from . import _array_module as xp, api_version
+from . import array_helpers as ah
 from . import dtype_helpers as dh
 from . import shape_helpers as sh
 from . import xps
@@ -211,6 +212,7 @@ def tuples(elements, *, min_size=0, max_size=None, unique_by=None, unique=False)
 
 # Use this to avoid memory errors with NumPy.
 # See https://github.com/numpy/numpy/issues/15753
+# Note, the hypothesis default for max_dims is min_dims + 2 (i.e., 0 + 2)
 def shapes(**kw):
     kw.setdefault('min_dims', 0)
     kw.setdefault('min_side', 0)
@@ -280,16 +282,19 @@ def mutually_broadcastable_shapes(
 
 # Note: This should become hermitian_matrices when complex dtypes are added
 @composite
-def symmetric_matrices(draw, dtypes=xps.floating_dtypes(), finite=True):
+def symmetric_matrices(draw, dtypes=xps.floating_dtypes(), finite=True, bound=10.):
     shape = draw(square_matrix_shapes)
     dtype = draw(dtypes)
     if not isinstance(finite, bool):
         finite = draw(finite)
     elements = {'allow_nan': False, 'allow_infinity': False} if finite else None
     a = draw(arrays(dtype=dtype, shape=shape, elements=elements))
-    upper = xp.triu(a)
-    lower = xp.triu(a, k=1).mT
-    return upper + lower
+    at = ah._matrix_transpose(a)
+    H = (a + at)*0.5
+    if finite:
+        assume(not xp.any(xp.isinf(H)))
+    assume(xp.all((H == 0.) | ((1/bound <= xp.abs(H)) & (xp.abs(H) <= bound))))
+    return H
 
 @composite
 def positive_definite_matrices(draw, dtypes=xps.floating_dtypes()):

array_api_tests/meta/test_linalg.py

@@ -0,0 +1,16 @@
+import pytest
+
+from hypothesis import given
+
+from ..hypothesis_helpers import symmetric_matrices
+from .. import array_helpers as ah
+from .. import _array_module as xp
+
+@pytest.mark.xp_extension('linalg')
+@given(x=symmetric_matrices(finite=True))
+def test_symmetric_matrices(x):
+    upper = xp.triu(x)
+    lower = xp.tril(x)
+    lowerT = ah._matrix_transpose(lower)
+
+    ah.assert_exactly_equal(upper, lowerT)