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=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.
@@ -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,23 @@
 # Standin strategy for not yet implemented tests
 todo = none()
 
+def assert_equal(x, y):
+ if x.dtype in dh.float_dtypes:
+ # It's too difficult to do an approximately equal test here because
+ # different routines can give completely different answers, and even
+ # when it does work, the elementwise comparisons are too slow. So for
+ # floating-point dtypes only test the shape and dtypes.
+
+ # assert_allclose(x, y)
+
+ 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})"
+ 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
 
@@ -225,7 +242,6 @@ def true_diag(x_stack):
 
  _test_stacks(linalg.diagonal, x, **kw, res=res, dims=1, true_val=true_diag)
 
-@pytest.mark.skip(reason="Inputs need to be restricted") # TODO
 @pytest.mark.xp_extension('linalg')
 @given(x=symmetric_matrices(finite=True))
 def test_eigh(x):
@@ -242,8 +258,15 @@ def test_eigh(x):
  assert eigenvectors.dtype == x.dtype, "eigh().eigenvectors did not return the correct dtype"
  assert eigenvectors.shape == x.shape, "eigh().eigenvectors did not return the correct shape"
 
+ # Note: _test_stacks here is only testing the shape and dtype. The actual
+ # eigenvalues and eigenvectors may not be equal at all, since there is not
+ # requirements about how eigh computes an eigenbasis, or about the order
+ # of the eigenvalues
  _test_stacks(lambda x: linalg.eigh(x).eigenvalues, x,
  res=eigenvalues, dims=1)
+
+ # TODO: Test that eigenvectors are orthonormal.
+
  _test_stacks(lambda x: linalg.eigh(x).eigenvectors, x,
  res=eigenvectors, dims=2)
 
@@ -258,9 +281,14 @@ def test_eigvalsh(x):
  assert res.dtype == x.dtype, "eigvalsh() did not return the correct dtype"
  assert res.shape == x.shape[:-1], "eigvalsh() did not return the correct shape"
 
+ # Note: _test_stacks here is only testing the shape and dtype. The actual
+ # eigenvalues may not be equal at all, since there is not requirements or
+ # about the order of the eigenvalues, and the stacking code may use a
+ # different code path.
  _test_stacks(linalg.eigvalsh, x, res=res, dims=1)
 
  # TODO: Should we test that the result is the same as eigh(x).eigenvalues?
+ # (probably no because the spec doesn't actually require that)
 
  # TODO: Test that res actually corresponds to the eigenvalues of x
 
@@ -309,8 +337,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 +597,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 +767,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
+ )