More improvements to test_linalg

.github/workflows/numpy.yml

@@ -37,6 +37,9 @@ jobs:
         # The return dtype for trace is not consistent in the spec
         # https://github.com/data-apis/array-api/issues/202#issuecomment-952529197
         array_api_tests/test_linalg.py::test_trace
+        # Various fixes to vector_norm are in
+        # https://github.com/numpy/numpy/pull/21084.
+        array_api_tests/test_linalg.py::test_vector_norm
         # waiting on NumPy to allow/revert distinct NaNs for np.unique
         # https://github.com/numpy/numpy/issues/20326#issuecomment-1012380448
         array_api_tests/test_set_functions.py

array_api_tests/array_helpers.py

@@ -9,6 +9,9 @@
 from ._array_module import logical_not, subtract, floor, ceil, where
 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 +149,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=0.25, abs_tol=1):
+    """
+    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.
@@ -305,4 +345,3 @@ 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"
-

array_api_tests/test_linalg.py

@@ -15,19 +15,22 @@
 
 import pytest
 from hypothesis import assume, given
-from hypothesis.strategies import (booleans, composite, none, tuples, integers,
-                                   shared, sampled_from, one_of, data, just)
+from hypothesis.strategies import (booleans, composite, none, tuples, floats,
+                                   integers, shared, sampled_from, one_of,
+                                   data, just)
 from ndindex import iter_indices
 
-from .array_helpers import assert_exactly_equal, asarray
+import itertools
+
+from .array_helpers import assert_exactly_equal, asarray, assert_allclose
 from .hypothesis_helpers import (xps, dtypes, shapes, kwargs, matrix_shapes,
                                  square_matrix_shapes, symmetric_matrices,
                                  positive_definite_matrices, MAX_ARRAY_SIZE,
                                  invertible_matrices, two_mutual_arrays,
                                  mutually_promotable_dtypes, one_d_shapes,
                                  two_mutually_broadcastable_shapes,
                                  SQRT_MAX_ARRAY_SIZE, finite_matrices,
-                                 rtol_shared_matrix_shapes, rtols)
+                                 rtol_shared_matrix_shapes, rtols, axes)
 from . import dtype_helpers as dh
 from . import pytest_helpers as ph
 from . import shape_helpers as sh
@@ -41,9 +44,15 @@
 # Standin strategy for not yet implemented tests
 todo = none()
 
+def assert_equal(x, y):
+    if x.dtype in dh.float_dtypes:
+        assert_allclose(x, y)
+    else:
+        assert_exactly_equal(x, y)
+
 def _test_stacks(f, *args, res=None, dims=2, true_val=None,
                  matrix_axes=(-2, -1),
-                 assert_equal=assert_exactly_equal, **kw):
+                 assert_equal=assert_equal, **kw):
     """
     Test that f(*args, **kw) maps across stacks of matrices
 
@@ -309,8 +318,6 @@ def test_matmul(x1, x2):
         assert res.shape == stack_shape + (x1.shape[-2], x2.shape[-1])
         _test_stacks(_array_module.matmul, x1, x2, res=res)
 
-matrix_norm_shapes = shared(matrix_shapes())
-
 @pytest.mark.xp_extension('linalg')
 @given(
     x=finite_matrices(),
@@ -571,22 +578,118 @@ def test_svdvals(x):
 
     # TODO: Check that svdvals() is the same as svd().s.
 
+_tensordot_pre_shapes = shared(two_mutually_broadcastable_shapes)
+
+@composite
+def _tensordot_axes(draw):
+    shape1, shape2 = draw(_tensordot_pre_shapes)
+    ndim1, ndim2 = len(shape1), len(shape2)
+    isint = draw(booleans())
+
+    if isint:
+        N = min(ndim1, ndim2)
+        return draw(integers(0, N))
+    else:
+        if ndim1 < ndim2:
+            first = draw(xps.valid_tuple_axes(ndim1))
+            second = draw(xps.valid_tuple_axes(ndim2, min_size=len(first),
+                                               max_size=len(first)))
+        else:
+            second = draw(xps.valid_tuple_axes(ndim2))
+            first = draw(xps.valid_tuple_axes(ndim1, min_size=len(second),
+                                               max_size=len(second)))
+        return (tuple(first), tuple(second))
+
+tensordot_kw = shared(kwargs(axes=_tensordot_axes()))
+
+@composite
+def tensordot_shapes(draw):
+    _shape1, _shape2 = map(list, draw(_tensordot_pre_shapes))
+    ndim1, ndim2 = len(_shape1), len(_shape2)
+    kw = draw(tensordot_kw)
+    if 'axes' not in kw:
+        assume(ndim1 >= 2 and ndim2 >= 2)
+    axes = kw.get('axes', 2)
+
+    if isinstance(axes, int):
+        axes = [list(range(-axes, 0)), list(range(0, axes))]
+
+    first, second = axes
+    for i, j in zip(first, second):
+        try:
+            if -ndim2 <= j < ndim2 and _shape2[j] != 1:
+                _shape1[i] = _shape2[j]
+            if -ndim1 <= i < ndim1 and _shape1[i] != 1:
+                _shape2[j] = _shape1[i]
+        except:
+            raise
+
+    shape1, shape2 = map(tuple, [_shape1, _shape2])
+    return (shape1, shape2)
+
+def _test_tensordot_stacks(x1, x2, kw, res):
+    """
+    Variant of _test_stacks for tensordot
+
+    tensordot doesn't stack directly along the non-contracted dimensions like
+    the other linalg functions. Rather, it is stacked along the product of
+    each non-contracted dimension. These dimensions are independent of one
+    another and do not broadcast.
+    """
+    shape1, shape2 = x1.shape, x2.shape
+
+    axes = kw.get('axes', 2)
+
+    if isinstance(axes, int):
+        res_axes = axes
+        axes = [list(range(-axes, 0)), list(range(0, axes))]
+    else:
+        # Convert something like (0, 4, 2) into (0, 2, 1)
+        res_axes = []
+        for a, s in zip(axes, [shape1, shape2]):
+            indices = [range(len(s))[i] for i in a]
+            repl = dict(zip(sorted(indices), range(len(indices))))
+            res_axes.append(tuple(repl[i] for i in indices))
+
+    for ((i,), (j,)), (res_idx,) in zip(
+            itertools.product(
+                iter_indices(shape1, skip_axes=axes[0]),
+                iter_indices(shape2, skip_axes=axes[1])),
+            iter_indices(res.shape)):
+        i, j, res_idx = i.raw, j.raw, res_idx.raw
+
+        res_stack = res[res_idx]
+        x1_stack = x1[i]
+        x2_stack = x2[j]
+        decomp_res_stack = xp.tensordot(x1_stack, x2_stack, axes=res_axes)
+        assert_exactly_equal(res_stack, decomp_res_stack)
 
 @given(
-    dtypes=mutually_promotable_dtypes(dtypes=dh.numeric_dtypes),
-    shape=shapes(),
-    data=data(),
+    *two_mutual_arrays(dh.numeric_dtypes, two_shapes=tensordot_shapes()),
+    tensordot_kw,
 )
-def test_tensordot(dtypes, shape, data):
+def test_tensordot(x1, x2, kw):
     # TODO: vary shapes, vary contracted axes, test different axes arguments
-    x1 = data.draw(xps.arrays(dtype=dtypes[0], shape=shape), label="x1")
-    x2 = data.draw(xps.arrays(dtype=dtypes[1], shape=shape), label="x2")
+    res = xp.tensordot(x1, x2, **kw)
 
-    out = xp.tensordot(x1, x2, axes=len(shape))
+    ph.assert_dtype("tensordot", [x1.dtype, x2.dtype], res.dtype)
 
-    ph.assert_dtype("tensordot", dtypes, out.dtype)
-    # TODO: assert shape and elements
+    axes = _axes = kw.get('axes', 2)
 
+    if isinstance(axes, int):
+        _axes = [list(range(-axes, 0)), list(range(0, axes))]
+
+    _shape1 = list(x1.shape)
+    _shape2 = list(x2.shape)
+    for i, j in zip(*_axes):
+        _shape1[i] = _shape2[j] = None
+    _shape1 = tuple([i for i in _shape1 if i is not None])
+    _shape2 = tuple([i for i in _shape2 if i is not None])
+    result_shape = _shape1 + _shape2
+    ph.assert_result_shape('tensordot', [x1.shape, x2.shape], res.shape,
+                           expected=result_shape)
+    # TODO: assert stacking and elements
+    _test_tensordot_stacks(x1, x2, kw, res)
 
 @pytest.mark.xp_extension('linalg')
 @given(
@@ -645,11 +748,42 @@ def test_vecdot(dtypes, shape, data):
     # TODO: assert shape and elements
 
 
+# Insanely large orders might not work. There isn't a limit specified in the
+# spec, so we just limit to reasonable values here.
+max_ord = 100
+
 @pytest.mark.xp_extension('linalg')
 @given(
-    x=xps.arrays(dtype=xps.floating_dtypes(), shape=shapes()),
-    kw=kwargs(axis=todo, keepdims=todo, ord=todo)
+    x=xps.arrays(dtype=xps.floating_dtypes(), shape=shapes(min_side=1)),
+    data=data(),
 )
-def test_vector_norm(x, kw):
-    # res = linalg.vector_norm(x, **kw)
-    pass
+def test_vector_norm(x, data):
+    kw = data.draw(
+        # We use data because axes is parameterized on x.ndim
+        kwargs(axis=axes(x.ndim),
+               keepdims=booleans(),
+               ord=one_of(
+                   sampled_from([2, 1, 0, -1, -2, float("inf"), float("-inf")]),
+                   integers(-max_ord, max_ord),
+                   floats(-max_ord, max_ord),
+               )), label="kw")
+
+
+    res = linalg.vector_norm(x, **kw)
+    axis = kw.get('axis', None)
+    keepdims = kw.get('keepdims', False)
+    # TODO: Check that the ord values give the correct norms.
+    # ord = kw.get('ord', 2)
+
+    _axes = sh.normalise_axis(axis, x.ndim)
+
+    ph.assert_keepdimable_shape('linalg.vector_norm', res.shape, x.shape,
+                                _axes, keepdims, **kw)
+    ph.assert_dtype('linalg.vector_norm', x.dtype, res.dtype)
+
+    _kw = kw.copy()
+    _kw.pop('axis', None)
+    _test_stacks(linalg.vector_norm, x, res=res,
+                 dims=x.ndim if keepdims else 0,
+                 matrix_axes=_axes, **_kw
+                 )