From 360df9fd378160fe9addc86239c77596f221d7be Mon Sep 17 00:00:00 2001
From: Aaron Meurer <asmeurer@gmail.com>
Date: Thu, 17 Feb 2022 15:39:48 -0700
Subject: [PATCH 01/54] Add shape testing for vector_norm()

---
 array_api_tests/test_linalg.py | 50 ++++++++++++++++++++++++++++------
 1 file changed, 42 insertions(+), 8 deletions(-)

diff --git a/array_api_tests/test_linalg.py b/array_api_tests/test_linalg.py
index be4a22ca..d6fde7ce 100644
--- a/array_api_tests/test_linalg.py
+++ b/array_api_tests/test_linalg.py
@@ -15,8 +15,9 @@
 
 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
@@ -27,7 +28,7 @@
                                  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
@@ -645,11 +646,44 @@ 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(
+                   just(1),
+                   just(2),
+                   just(float('inf')),
+                   integers(1, max_ord),
+                   floats(1, max_ord, allow_nan=False,
+                          allow_infinity=False),
+                   just(0),
+                   just(-1),
+                   just(-2),
+                   just(float('-inf')),
+                   integers(-max_ord, -1),
+                   floats(-max_ord, -1, allow_nan=False,
+                          allow_infinity=False),
+               )), 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)

From e34f492e54f90bed25a19b229129e2d4715a36e0 Mon Sep 17 00:00:00 2001
From: Aaron Meurer <asmeurer@gmail.com>
Date: Thu, 17 Feb 2022 18:39:38 -0700
Subject: [PATCH 02/54] Test the dtype and stacks in the vector_norm() test

---
 array_api_tests/test_linalg.py | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/array_api_tests/test_linalg.py b/array_api_tests/test_linalg.py
index d6fde7ce..3b98c1b6 100644
--- a/array_api_tests/test_linalg.py
+++ b/array_api_tests/test_linalg.py
@@ -687,3 +687,11 @@ def test_vector_norm(x, data):
 
     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
+                 )

From 111c2372e8d4e16ad8f71ae5b196b4aceccd5d9a Mon Sep 17 00:00:00 2001
From: Aaron Meurer <asmeurer@gmail.com>
Date: Thu, 17 Feb 2022 18:40:10 -0700
Subject: [PATCH 03/54] Remove an ununsed variable

---
 array_api_tests/test_linalg.py | 2 --
 1 file changed, 2 deletions(-)

diff --git a/array_api_tests/test_linalg.py b/array_api_tests/test_linalg.py
index 3b98c1b6..4fc4202a 100644
--- a/array_api_tests/test_linalg.py
+++ b/array_api_tests/test_linalg.py
@@ -310,8 +310,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(),

From 4f3aa54d05c8654d7e4971c9beef8466215b4870 Mon Sep 17 00:00:00 2001
From: Aaron Meurer <asmeurer@gmail.com>
Date: Fri, 18 Feb 2022 19:11:47 -0600
Subject: [PATCH 04/54] Use a simpler strategy for ord in test_vector_norm

Co-authored-by: Matthew Barber <quitesimplymatt@gmail.com>
---
 array_api_tests/test_linalg.py | 16 +++-------------
 1 file changed, 3 insertions(+), 13 deletions(-)

diff --git a/array_api_tests/test_linalg.py b/array_api_tests/test_linalg.py
index 4fc4202a..04b6b0fa 100644
--- a/array_api_tests/test_linalg.py
+++ b/array_api_tests/test_linalg.py
@@ -659,19 +659,9 @@ def test_vector_norm(x, data):
         kwargs(axis=axes(x.ndim),
                keepdims=booleans(),
                ord=one_of(
-                   just(1),
-                   just(2),
-                   just(float('inf')),
-                   integers(1, max_ord),
-                   floats(1, max_ord, allow_nan=False,
-                          allow_infinity=False),
-                   just(0),
-                   just(-1),
-                   just(-2),
-                   just(float('-inf')),
-                   integers(-max_ord, -1),
-                   floats(-max_ord, -1, allow_nan=False,
-                          allow_infinity=False),
+                   sampled_from([2, 1, 0, -1, -2, float("inf"), float("-inf")],
+                   integers(-max_ord, max_ord),
+                   floats(-max_ord, max_ord),
                )), label="kw")
 
 

From 979b81b1adb52fc1c3b083f2b7d25cbc8028f93a Mon Sep 17 00:00:00 2001
From: Aaron Meurer <asmeurer@gmail.com>
Date: Fri, 25 Feb 2022 16:45:35 -0700
Subject: [PATCH 05/54] Skip the test_vector_norm test on the NumPy CI

The upstream has bugs that will be fixed in the next release.
---
 .github/workflows/numpy.yml | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/.github/workflows/numpy.yml b/.github/workflows/numpy.yml
index 82420452..feef11c8 100644
--- a/.github/workflows/numpy.yml
+++ b/.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

From 8df237af02357c2be4d880e59e9768866ce12cb3 Mon Sep 17 00:00:00 2001
From: Aaron Meurer <asmeurer@gmail.com>
Date: Fri, 25 Feb 2022 16:46:31 -0700
Subject: [PATCH 06/54] Fix syntax error

---
 array_api_tests/test_linalg.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/array_api_tests/test_linalg.py b/array_api_tests/test_linalg.py
index 04b6b0fa..f00e3319 100644
--- a/array_api_tests/test_linalg.py
+++ b/array_api_tests/test_linalg.py
@@ -659,7 +659,7 @@ def test_vector_norm(x, data):
         kwargs(axis=axes(x.ndim),
                keepdims=booleans(),
                ord=one_of(
-                   sampled_from([2, 1, 0, -1, -2, float("inf"), float("-inf")],
+                   sampled_from([2, 1, 0, -1, -2, float("inf"), float("-inf")]),
                    integers(-max_ord, max_ord),
                    floats(-max_ord, max_ord),
                )), label="kw")

From d11a685b8111412fa792eb0432349df86a675d24 Mon Sep 17 00:00:00 2001
From: Aaron Meurer <asmeurer@gmail.com>
Date: Fri, 25 Feb 2022 19:42:03 -0700
Subject: [PATCH 07/54] Fix the input strategies for test_tensordot()

---
 array_api_tests/test_linalg.py | 68 ++++++++++++++++++++++++++++------
 1 file changed, 56 insertions(+), 12 deletions(-)

diff --git a/array_api_tests/test_linalg.py b/array_api_tests/test_linalg.py
index f00e3319..82a1ed71 100644
--- a/array_api_tests/test_linalg.py
+++ b/array_api_tests/test_linalg.py
@@ -15,9 +15,9 @@
 
 import pytest
 from hypothesis import assume, given
-from hypothesis.strategies import (booleans, composite, none, tuples, floats,
-                                   integers, shared, sampled_from, one_of,
-                                   data, just)
+from hypothesis.strategies import (booleans, composite, none, lists, tuples,
+                                   floats, integers, shared, sampled_from,
+                                   one_of, data, just)
 from ndindex import iter_indices
 
 from .array_helpers import assert_exactly_equal, asarray
@@ -570,20 +570,64 @@ 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)
 
 @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")
-
-    out = xp.tensordot(x1, x2, axes=len(shape))
+    out = xp.tensordot(x1, x2, **kw)
 
-    ph.assert_dtype("tensordot", dtypes, out.dtype)
+    ph.assert_dtype("tensordot", [x1.dtype, x2.dtype], out.dtype)
     # TODO: assert shape and elements
 
 

From a776cd43f614d2821482ac7d9a20a3ce2dc4a1a9 Mon Sep 17 00:00:00 2001
From: Aaron Meurer <asmeurer@gmail.com>
Date: Sat, 9 Apr 2022 00:07:24 -0600
Subject: [PATCH 08/54] Add a test for the tensordot result shape

---
 array_api_tests/test_linalg.py | 17 ++++++++++++++++-
 1 file changed, 16 insertions(+), 1 deletion(-)

diff --git a/array_api_tests/test_linalg.py b/array_api_tests/test_linalg.py
index 82a1ed71..45aab88f 100644
--- a/array_api_tests/test_linalg.py
+++ b/array_api_tests/test_linalg.py
@@ -628,7 +628,22 @@ def test_tensordot(x1, x2, kw):
     out = xp.tensordot(x1, x2, **kw)
 
     ph.assert_dtype("tensordot", [x1.dtype, x2.dtype], 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], out.shape,
+                           expected=result_shape)
+    # TODO: assert stacking and elements
 
 
 @pytest.mark.xp_extension('linalg')

From 45b36d600b2e63bb08ae39799db10b598785b21e Mon Sep 17 00:00:00 2001
From: Aaron Meurer <asmeurer@gmail.com>
Date: Tue, 12 Apr 2022 01:10:46 -0600
Subject: [PATCH 09/54] Test stacking for tensordot

---
 array_api_tests/test_linalg.py | 53 +++++++++++++++++++++++++++++-----
 1 file changed, 46 insertions(+), 7 deletions(-)

diff --git a/array_api_tests/test_linalg.py b/array_api_tests/test_linalg.py
index 45aab88f..ae37196a 100644
--- a/array_api_tests/test_linalg.py
+++ b/array_api_tests/test_linalg.py
@@ -15,11 +15,13 @@
 
 import pytest
 from hypothesis import assume, given
-from hypothesis.strategies import (booleans, composite, none, lists, tuples,
-                                   floats, 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
 
+import itertools
+
 from .array_helpers import assert_exactly_equal, asarray
 from .hypothesis_helpers import (xps, dtypes, shapes, kwargs, matrix_shapes,
                                  square_matrix_shapes, symmetric_matrices,
@@ -619,15 +621,52 @@ def tensordot_shapes(draw):
     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(
     *two_mutual_arrays(dh.numeric_dtypes, two_shapes=tensordot_shapes()),
     tensordot_kw,
 )
 def test_tensordot(x1, x2, kw):
     # TODO: vary shapes, vary contracted axes, test different axes arguments
-    out = xp.tensordot(x1, x2, **kw)
+    res = xp.tensordot(x1, x2, **kw)
 
-    ph.assert_dtype("tensordot", [x1.dtype, x2.dtype], out.dtype)
+    ph.assert_dtype("tensordot", [x1.dtype, x2.dtype], res.dtype)
 
     axes = _axes = kw.get('axes', 2)
 
@@ -641,10 +680,10 @@ def test_tensordot(x1, x2, kw):
     _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], out.shape,
+    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(

From 414b3223cdd348f0cc891d363d39d749450e7b23 Mon Sep 17 00:00:00 2001
From: Aaron Meurer <asmeurer@gmail.com>
Date: Mon, 25 Apr 2022 17:20:43 -0600
Subject: [PATCH 10/54] Add allclose() and assert_allclose() helper functions

---
 array_api_tests/array_helpers.py | 41 +++++++++++++++++++++++++++++++-
 1 file changed, 40 insertions(+), 1 deletion(-)

diff --git a/array_api_tests/array_helpers.py b/array_api_tests/array_helpers.py
index ef4f719a..0ec73ba2 100644
--- a/array_api_tests/array_helpers.py
+++ b/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"
-

From 9bb8c7ac4ec256856db663dc1b31e19ba3092d2a Mon Sep 17 00:00:00 2001
From: Aaron Meurer <asmeurer@gmail.com>
Date: Mon, 25 Apr 2022 17:21:01 -0600
Subject: [PATCH 11/54] Use assert_allclose() in the linalg tests for float
 inputs

Fixes #44.
Related to #117.
---
 array_api_tests/test_linalg.py | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/array_api_tests/test_linalg.py b/array_api_tests/test_linalg.py
index ae37196a..64562e7d 100644
--- a/array_api_tests/test_linalg.py
+++ b/array_api_tests/test_linalg.py
@@ -22,7 +22,7 @@
 
 import itertools
 
-from .array_helpers import assert_exactly_equal, asarray
+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,
@@ -44,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
 

From b3fb4ecafa759cc125b722532bc66094c0e77cc3 Mon Sep 17 00:00:00 2001
From: Aaron Meurer <asmeurer@gmail.com>
Date: Tue, 26 Apr 2022 14:24:28 -0600
Subject: [PATCH 12/54] Remove skip from test_eigh

---
 array_api_tests/test_linalg.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/array_api_tests/test_linalg.py b/array_api_tests/test_linalg.py
index 64562e7d..19b88e34 100644
--- a/array_api_tests/test_linalg.py
+++ b/array_api_tests/test_linalg.py
@@ -234,7 +234,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):

From 241220e386009afd4cab657f83a54f8d013159a3 Mon Sep 17 00:00:00 2001
From: Aaron Meurer <asmeurer@gmail.com>
Date: Thu, 5 May 2022 19:30:00 -0600
Subject: [PATCH 13/54] Disable eigenvectors stack test

There are different equivalent ways of representing eigenspaces in general,
and the same algorithm may choose different ways for stacked vs. non-stacked
matrices (e.g., they differ for cupy).
---
 array_api_tests/test_linalg.py | 9 +++++++--
 1 file changed, 7 insertions(+), 2 deletions(-)

diff --git a/array_api_tests/test_linalg.py b/array_api_tests/test_linalg.py
index 64562e7d..b1ad2510 100644
--- a/array_api_tests/test_linalg.py
+++ b/array_api_tests/test_linalg.py
@@ -253,8 +253,13 @@ def test_eigh(x):
 
     _test_stacks(lambda x: linalg.eigh(x).eigenvalues, x,
                  res=eigenvalues, dims=1)
-    _test_stacks(lambda x: linalg.eigh(x).eigenvectors, x,
-                 res=eigenvectors, dims=2)
+
+    # There are equivalent ways of representing eigenvectors, and algorithms
+    # may not give the same eigenvectors on a stack vs. a matrix.
+    # TODO: Test that eigenvectors are orthonormal.
+
+    # _test_stacks(lambda x: linalg.eigh(x).eigenvectors, x,
+    #              res=eigenvectors, dims=2)
 
     # TODO: Test that res actually corresponds to the eigenvalues and
     # eigenvectors of x

From ca70fbe3046f91fc9286ed1270b049e7fde7169e Mon Sep 17 00:00:00 2001
From: Aaron Meurer <asmeurer@gmail.com>
Date: Thu, 5 May 2022 23:47:41 -0600
Subject: [PATCH 14/54] Reduce the relative tolerance in assert_allclose

This still is going to require some tweaking.
---
 array_api_tests/array_helpers.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/array_api_tests/array_helpers.py b/array_api_tests/array_helpers.py
index 0ec73ba2..701b3e3f 100644
--- a/array_api_tests/array_helpers.py
+++ b/array_api_tests/array_helpers.py
@@ -171,7 +171,7 @@ def allclose(x, y, rel_tol=0.25, abs_tol=1, return_indices=False):
             return False
     return True
 
-def assert_allclose(x, y, rel_tol=0.25, abs_tol=1):
+def assert_allclose(x, y, rel_tol=1, abs_tol=0.):
     """
     Test that x and y are approximately equal to each other.
 

From 720b3091e6d8e4152f86b990729d66fc0772b403 Mon Sep 17 00:00:00 2001
From: Aaron Meurer <asmeurer@gmail.com>
Date: Thu, 5 May 2022 23:48:05 -0600
Subject: [PATCH 15/54] Sort the eigenvalues when testing stacks

---
 array_api_tests/test_linalg.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/array_api_tests/test_linalg.py b/array_api_tests/test_linalg.py
index b1ad2510..a551e876 100644
--- a/array_api_tests/test_linalg.py
+++ b/array_api_tests/test_linalg.py
@@ -251,8 +251,11 @@ 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"
 
+    # The order of the eigenvectors is not specified, so make sure the same
+    # eigenvectors are compared against each other.
     _test_stacks(lambda x: linalg.eigh(x).eigenvalues, x,
-                 res=eigenvalues, dims=1)
+                 res=eigenvalues, dims=1, assert_equal=lambda a, b:
+                 assert_equal(xp.sort(a), xp.sort(b)))
 
     # There are equivalent ways of representing eigenvectors, and algorithms
     # may not give the same eigenvectors on a stack vs. a matrix.

From f4392598017b2e0cda4022d377c11a166baddc9a Mon Sep 17 00:00:00 2001
From: Aaron Meurer <asmeurer@gmail.com>
Date: Fri, 3 Jun 2022 15:13:48 -0600
Subject: [PATCH 16/54] Sort the results in eigvalsh before comparing

---
 array_api_tests/test_linalg.py | 9 ++++++---
 1 file changed, 6 insertions(+), 3 deletions(-)

diff --git a/array_api_tests/test_linalg.py b/array_api_tests/test_linalg.py
index 4ec4d243..702c0a81 100644
--- a/array_api_tests/test_linalg.py
+++ b/array_api_tests/test_linalg.py
@@ -250,8 +250,8 @@ 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"
 
-    # The order of the eigenvectors is not specified, so make sure the same
-    # eigenvectors are compared against each other.
+    # The order of the eigenvalues is not specified, so make sure the same
+    # eigenvalues are compared against each other.
     _test_stacks(lambda x: linalg.eigh(x).eigenvalues, x,
                  res=eigenvalues, dims=1, assert_equal=lambda a, b:
                  assert_equal(xp.sort(a), xp.sort(b)))
@@ -274,7 +274,10 @@ 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"
 
-    _test_stacks(linalg.eigvalsh, x, res=res, dims=1)
+    # The order of the eigenvalues is not specified, so make sure the same
+    # eigenvalues are compared against each other.
+    _test_stacks(linalg.eigvalsh, x, res=res, dims=1, assert_equal=lambda a,
+                 b: assert_equal(xp.sort(a), xp.sort(b)))
 
     # TODO: Should we test that the result is the same as eigh(x).eigenvalues?
 

From 75ca73af786baf9f6d59c8dc3888ee116651f2e3 Mon Sep 17 00:00:00 2001
From: Aaron Meurer <asmeurer@gmail.com>
Date: Sun, 12 Jun 2022 18:50:56 -0600
Subject: [PATCH 17/54] Remove the allclose testing in linalg

Floating-point dtypes now only test shape and dtype in the stacking test.
Trying to make the allclose test work is too difficult for linear algebra
functions.
---
 array_api_tests/test_linalg.py | 35 +++++++++++++++++++++-------------
 1 file changed, 22 insertions(+), 13 deletions(-)

diff --git a/array_api_tests/test_linalg.py b/array_api_tests/test_linalg.py
index 702c0a81..b8b9724e 100644
--- a/array_api_tests/test_linalg.py
+++ b/array_api_tests/test_linalg.py
@@ -46,7 +46,15 @@
 
 def assert_equal(x, y):
     if x.dtype in dh.float_dtypes:
-        assert_allclose(x, y)
+        # 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)
 
@@ -250,18 +258,17 @@ 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"
 
-    # The order of the eigenvalues is not specified, so make sure the same
-    # eigenvalues are compared against each other.
+    # 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, assert_equal=lambda a, b:
-                 assert_equal(xp.sort(a), xp.sort(b)))
+                 res=eigenvalues, dims=1)
 
-    # There are equivalent ways of representing eigenvectors, and algorithms
-    # may not give the same eigenvectors on a stack vs. a matrix.
     # TODO: Test that eigenvectors are orthonormal.
 
-    # _test_stacks(lambda x: linalg.eigh(x).eigenvectors, x,
-    #              res=eigenvectors, dims=2)
+    _test_stacks(lambda x: linalg.eigh(x).eigenvectors, x,
+                 res=eigenvectors, dims=2)
 
     # TODO: Test that res actually corresponds to the eigenvalues and
     # eigenvectors of x
@@ -274,12 +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"
 
-    # The order of the eigenvalues is not specified, so make sure the same
-    # eigenvalues are compared against each other.
-    _test_stacks(linalg.eigvalsh, x, res=res, dims=1, assert_equal=lambda a,
-                 b: assert_equal(xp.sort(a), xp.sort(b)))
+    # 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
 

From d86a0a1c98095b9ebd08878991934a575cfcea65 Mon Sep 17 00:00:00 2001
From: Aaron Meurer <asmeurer@gmail.com>
Date: Thu, 16 Jun 2022 15:28:47 -0600
Subject: [PATCH 18/54] Add (commented out) stacking tests for solve()

---
 array_api_tests/test_linalg.py | 10 +++++++++-
 1 file changed, 9 insertions(+), 1 deletion(-)

diff --git a/array_api_tests/test_linalg.py b/array_api_tests/test_linalg.py
index b8b9724e..c3cfbf11 100644
--- a/array_api_tests/test_linalg.py
+++ b/array_api_tests/test_linalg.py
@@ -539,7 +539,15 @@ def _x2_shapes(draw):
 @pytest.mark.xp_extension('linalg')
 @given(*solve_args())
 def test_solve(x1, x2):
-    linalg.solve(x1, x2)
+    res = linalg.solve(x1, x2)
+
+    # TODO: This requires an upstream fix to ndindex
+    # (https://github.com/Quansight-Labs/ndindex/pull/131)
+
+    # if x2.ndim == 1:
+    #     _test_stacks(linalg.solve, x1, x2, res=res, dims=1)
+    # else:
+    #     _test_stacks(linalg.solve, x1, x2, res=res, dims=2)
 
 @pytest.mark.xp_extension('linalg')
 @given(

From 9bccfa5c6c488fa5e76e8a72c0899ff92522f46d Mon Sep 17 00:00:00 2001
From: Aaron Meurer <asmeurer@gmail.com>
Date: Thu, 16 Jun 2022 16:11:10 -0600
Subject: [PATCH 19/54] Remove unused none standin in the linalg tests

---
 array_api_tests/test_linalg.py | 7 ++-----
 1 file changed, 2 insertions(+), 5 deletions(-)

diff --git a/array_api_tests/test_linalg.py b/array_api_tests/test_linalg.py
index c3cfbf11..d1629af8 100644
--- a/array_api_tests/test_linalg.py
+++ b/array_api_tests/test_linalg.py
@@ -15,14 +15,14 @@
 
 import pytest
 from hypothesis import assume, given
-from hypothesis.strategies import (booleans, composite, none, tuples, floats,
+from hypothesis.strategies import (booleans, composite, tuples, floats,
                                    integers, shared, sampled_from, one_of,
                                    data, just)
 from ndindex import iter_indices
 
 import itertools
 
-from .array_helpers import assert_exactly_equal, asarray, assert_allclose
+from .array_helpers import assert_exactly_equal, asarray
 from .hypothesis_helpers import (xps, dtypes, shapes, kwargs, matrix_shapes,
                                  square_matrix_shapes, symmetric_matrices,
                                  positive_definite_matrices, MAX_ARRAY_SIZE,
@@ -41,9 +41,6 @@
 
 pytestmark = pytest.mark.ci
 
-# 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

From f494b45244a3900bc63a6324ea2c9df569183b4d Mon Sep 17 00:00:00 2001
From: Aaron Meurer <asmeurer@gmail.com>
Date: Thu, 16 Jun 2022 16:17:33 -0600
Subject: [PATCH 20/54] Don't compare float elements in test_tensordot

---
 array_api_tests/test_linalg.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/array_api_tests/test_linalg.py b/array_api_tests/test_linalg.py
index d1629af8..aeb06e65 100644
--- a/array_api_tests/test_linalg.py
+++ b/array_api_tests/test_linalg.py
@@ -686,7 +686,7 @@ def _test_tensordot_stacks(x1, x2, kw, res):
         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)
+        assert_equal(res_stack, decomp_res_stack)
 
 @given(
     *two_mutual_arrays(dh.numeric_dtypes, two_shapes=tensordot_shapes()),
@@ -771,7 +771,6 @@ def test_vecdot(dtypes, shape, data):
     ph.assert_dtype("vecdot", dtypes, out.dtype)
     # 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

From 74add08b34d442d1cdde70c277120f03ba48ad33 Mon Sep 17 00:00:00 2001
From: Aaron Meurer <asmeurer@gmail.com>
Date: Fri, 24 Jun 2022 15:37:07 -0600
Subject: [PATCH 21/54] Fix test_vecdot

The NumPy implementation is currently incorrect, so I am not 100% if this test
is completely correct.
---
 array_api_tests/test_linalg.py | 32 +++++++++++++++++++++++---------
 1 file changed, 23 insertions(+), 9 deletions(-)

diff --git a/array_api_tests/test_linalg.py b/array_api_tests/test_linalg.py
index aeb06e65..4a13744d 100644
--- a/array_api_tests/test_linalg.py
+++ b/array_api_tests/test_linalg.py
@@ -17,7 +17,7 @@
 from hypothesis import assume, given
 from hypothesis.strategies import (booleans, composite, tuples, floats,
                                    integers, shared, sampled_from, one_of,
-                                   data, just)
+                                   data)
 from ndindex import iter_indices
 
 import itertools
@@ -29,6 +29,7 @@
                                  invertible_matrices, two_mutual_arrays,
                                  mutually_promotable_dtypes, one_d_shapes,
                                  two_mutually_broadcastable_shapes,
+                                 mutually_broadcastable_shapes,
                                  SQRT_MAX_ARRAY_SIZE, finite_matrices,
                                  rtol_shared_matrix_shapes, rtols, axes)
 from . import dtype_helpers as dh
@@ -756,20 +757,33 @@ def true_trace(x_stack):
 
 
 @given(
-    dtypes=mutually_promotable_dtypes(dtypes=dh.numeric_dtypes),
-    shape=shapes(min_dims=1),
-    data=data(),
+    *two_mutual_arrays(dh.numeric_dtypes, mutually_broadcastable_shapes(2, min_dims=1)),
+    kwargs(axis=integers()),
 )
-def test_vecdot(dtypes, shape, data):
+def test_vecdot(x1, x2, kw):
     # TODO: vary shapes, test different axis 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")
-    kw = data.draw(kwargs(axis=just(-1)))
+    broadcasted_shape = sh.broadcast_shapes(x1.shape, x2.shape)
+    ndim = len(broadcasted_shape)
+    axis = kw.get('axis', -1)
+    if not (-ndim <= axis < ndim):
+        ph.raises(Exception, lambda: xp.vecdot(x1, x2, **kw),
+                  f"vecdot did not raise an exception for invalid axis ({ndim=}, {kw=})")
+        return
+    x1_shape = (1,)*(ndim - x1.ndim) + tuple(x1.shape)
+    x2_shape = (1,)*(ndim - x1.ndim) + tuple(x2.shape)
+    if x1_shape[axis] != x2_shape[axis]:
+        ph.raises(Exception, lambda: xp.vecdot(x1, x2, **kw),
+                  "vecdot did not raise an exception for invalid shapes")
+        return
+    expected_shape = list(broadcasted_shape)
+    expected_shape.pop(axis)
+    expected_shape = tuple(expected_shape)
 
     out = xp.vecdot(x1, x2, **kw)
 
-    ph.assert_dtype("vecdot", dtypes, out.dtype)
+    ph.assert_dtype("vecdot", [x1.dtype, x2.dtype], out.dtype)
     # TODO: assert shape and elements
+    ph.assert_shape("vecdot", out.shape, expected_shape)
 
 # Insanely large orders might not work. There isn't a limit specified in the
 # spec, so we just limit to reasonable values here.

From f12be47678deca748f7223c57c73c90b80eb90e8 Mon Sep 17 00:00:00 2001
From: Aaron Meurer <asmeurer@gmail.com>
Date: Tue, 5 Jul 2022 17:42:29 -0600
Subject: [PATCH 22/54] Fix typo in test_vecdot

---
 array_api_tests/test_linalg.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/array_api_tests/test_linalg.py b/array_api_tests/test_linalg.py
index 4a13744d..33f45bc9 100644
--- a/array_api_tests/test_linalg.py
+++ b/array_api_tests/test_linalg.py
@@ -770,7 +770,7 @@ def test_vecdot(x1, x2, kw):
                   f"vecdot did not raise an exception for invalid axis ({ndim=}, {kw=})")
         return
     x1_shape = (1,)*(ndim - x1.ndim) + tuple(x1.shape)
-    x2_shape = (1,)*(ndim - x1.ndim) + tuple(x2.shape)
+    x2_shape = (1,)*(ndim - x2.ndim) + tuple(x2.shape)
     if x1_shape[axis] != x2_shape[axis]:
         ph.raises(Exception, lambda: xp.vecdot(x1, x2, **kw),
                   "vecdot did not raise an exception for invalid shapes")

From d41d0bd337a13c9a2fe62728a59da29e8dd2bf0c Mon Sep 17 00:00:00 2001
From: Aaron Meurer <asmeurer@gmail.com>
Date: Tue, 5 Jul 2022 17:42:37 -0600
Subject: [PATCH 23/54] Expand vecdot tests

---
 array_api_tests/test_linalg.py | 15 ++++++++++++---
 1 file changed, 12 insertions(+), 3 deletions(-)

diff --git a/array_api_tests/test_linalg.py b/array_api_tests/test_linalg.py
index 33f45bc9..0b2af5fc 100644
--- a/array_api_tests/test_linalg.py
+++ b/array_api_tests/test_linalg.py
@@ -779,11 +779,20 @@ def test_vecdot(x1, x2, kw):
     expected_shape.pop(axis)
     expected_shape = tuple(expected_shape)
 
-    out = xp.vecdot(x1, x2, **kw)
+    res = xp.vecdot(x1, x2, **kw)
 
-    ph.assert_dtype("vecdot", [x1.dtype, x2.dtype], out.dtype)
+    ph.assert_dtype("vecdot", [x1.dtype, x2.dtype], res.dtype)
     # TODO: assert shape and elements
-    ph.assert_shape("vecdot", out.shape, expected_shape)
+    ph.assert_shape("vecdot", res.shape, expected_shape)
+
+    if x1.dtype in dh.int_dtypes:
+        def true_val(x, y, axix=-1):
+            return xp.sum(x*y, dtype=res.dtype)
+    else:
+        true_val = None
+
+    _test_stacks(linalg.vecdot, x1, x2, res=res, dims=0,
+                 matrix_axes=(axis,), true_val=true_val)
 
 # Insanely large orders might not work. There isn't a limit specified in the
 # spec, so we just limit to reasonable values here.

From 48a8442f72877424a3a800df07b67d7b53798902 Mon Sep 17 00:00:00 2001
From: Aaron Meurer <asmeurer@gmail.com>
Date: Mon, 28 Nov 2022 22:54:59 -0700
Subject: [PATCH 24/54] Check specially that the result of linalg functions is
 not a unnamed tuple

---
 array_api_tests/test_linalg.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/array_api_tests/test_linalg.py b/array_api_tests/test_linalg.py
index 4c791575..78b1be30 100644
--- a/array_api_tests/test_linalg.py
+++ b/array_api_tests/test_linalg.py
@@ -106,6 +106,7 @@ def _test_namedtuple(res, fields, func_name):
     # a tuple subclass with the right fields in the right order.
 
     assert isinstance(res, tuple), f"{func_name}() did not return a tuple"
+    assert type(res) != tuple, f"{func_name}() did not return a namedtuple"
     assert len(res) == len(fields), f"{func_name}() result tuple not the correct length (should have {len(fields)} elements)"
     for i, field in enumerate(fields):
         assert hasattr(res, field), f"{func_name}() result namedtuple doesn't have the '{field}' field"

From fd6367fc9ad740e7204a2bba25a6430a7a42fbf3 Mon Sep 17 00:00:00 2001
From: Aaron Meurer <asmeurer@gmail.com>
Date: Fri, 17 Mar 2023 17:14:47 -0600
Subject: [PATCH 25/54] Use a more robust fallback helper for matrix_transpose

---
 array_api_tests/array_helpers.py      | 12 ++++++++++++
 array_api_tests/hypothesis_helpers.py |  4 +++-
 2 files changed, 15 insertions(+), 1 deletion(-)

diff --git a/array_api_tests/array_helpers.py b/array_api_tests/array_helpers.py
index 701b3e3f..51aa7bc8 100644
--- a/array_api_tests/array_helpers.py
+++ b/array_api_tests/array_helpers.py
@@ -7,6 +7,7 @@
 # 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
@@ -345,3 +346,14 @@ 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")
diff --git a/array_api_tests/hypothesis_helpers.py b/array_api_tests/hypothesis_helpers.py
index 20cc0e03..f4502c06 100644
--- a/array_api_tests/hypothesis_helpers.py
+++ b/array_api_tests/hypothesis_helpers.py
@@ -10,6 +10,7 @@
                                    sampled_from, shared)
 
 from . import _array_module as xp
+from . import array_helpers as ah
 from . import dtype_helpers as dh
 from . import shape_helpers as sh
 from . import xps
@@ -212,7 +213,8 @@ def symmetric_matrices(draw, dtypes=xps.floating_dtypes(), finite=True):
     elements = {'allow_nan': False, 'allow_infinity': False} if finite else None
     a = draw(xps.arrays(dtype=dtype, shape=shape, elements=elements))
     upper = xp.triu(a)
-    lower = xp.triu(a, k=1).mT
+    # mT and matrix_transpose are more likely to not be implemented
+    lower = ah._matrix_transpose(xp.triu(a, k=1))
     return upper + lower
 
 @composite

From 70177973a1c0d5df64a69d61635e44fa750393e5 Mon Sep 17 00:00:00 2001
From: Aaron Meurer <asmeurer@gmail.com>
Date: Mon, 20 Mar 2023 17:37:58 -0600
Subject: [PATCH 26/54] Be more constrained about constructing symmetric
 matrices

I think this still might be able to construct an ill-conditioned matrix, but
it didn't come up with any after thousands of example runs, let's keep it for
now.

Also construct symmetric matrices using (a + a.T)/2 instead of using
triu/tril, and, add a meta-test for it.
---
 array_api_tests/hypothesis_helpers.py | 12 +++++++-----
 array_api_tests/meta/test_linalg.py   | 16 ++++++++++++++++
 2 files changed, 23 insertions(+), 5 deletions(-)
 create mode 100644 array_api_tests/meta/test_linalg.py

diff --git a/array_api_tests/hypothesis_helpers.py b/array_api_tests/hypothesis_helpers.py
index f4502c06..a5c06ed9 100644
--- a/array_api_tests/hypothesis_helpers.py
+++ b/array_api_tests/hypothesis_helpers.py
@@ -207,15 +207,17 @@ 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)
     elements = {'allow_nan': False, 'allow_infinity': False} if finite else None
     a = draw(xps.arrays(dtype=dtype, shape=shape, elements=elements))
-    upper = xp.triu(a)
-    # mT and matrix_transpose are more likely to not be implemented
-    lower = ah._matrix_transpose(xp.triu(a, k=1))
-    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()):
diff --git a/array_api_tests/meta/test_linalg.py b/array_api_tests/meta/test_linalg.py
new file mode 100644
index 00000000..a4171e81
--- /dev/null
+++ b/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)

From 246e38a200c9f556186c35359e16b73ac16a5397 Mon Sep 17 00:00:00 2001
From: Aaron Meurer <asmeurer@gmail.com>
Date: Thu, 23 Mar 2023 17:06:35 -0600
Subject: [PATCH 27/54] Don't require the arguments to assert_keepdimable_shape
 to be positional-only

---
 array_api_tests/pytest_helpers.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/array_api_tests/pytest_helpers.py b/array_api_tests/pytest_helpers.py
index 051a063f..29496443 100644
--- a/array_api_tests/pytest_helpers.py
+++ b/array_api_tests/pytest_helpers.py
@@ -267,7 +267,6 @@ def assert_keepdimable_shape(
     out_shape: Shape,
     axes: Tuple[int, ...],
     keepdims: bool,
-    /,
     **kw,
 ):
     """

From 02542ff608164f571400d1421df9315b8f69512c Mon Sep 17 00:00:00 2001
From: Aaron Meurer <asmeurer@gmail.com>
Date: Tue, 28 Mar 2023 19:45:15 -0600
Subject: [PATCH 28/54] Show the arrays in the error message for
 assert_exactly_equal

---
 array_api_tests/array_helpers.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/array_api_tests/array_helpers.py b/array_api_tests/array_helpers.py
index 51aa7bc8..c34b87fb 100644
--- a/array_api_tests/array_helpers.py
+++ b/array_api_tests/array_helpers.py
@@ -217,7 +217,7 @@ def assert_exactly_equal(x, y):
 
     assert x.dtype == y.dtype, f"The input arrays do not have the same dtype ({x.dtype} != {y.dtype})"
 
-    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})"
 
 def assert_finite(x):
     """

From 72974e0ded4c8ea0f8e04bbe3aa0cb1246e1dc70 Mon Sep 17 00:00:00 2001
From: Aaron Meurer <asmeurer@gmail.com>
Date: Tue, 28 Mar 2023 19:46:11 -0600
Subject: [PATCH 29/54] Allow passing an extra assertion message to
 assert_equal in linalg and assert_exactly_equal

---
 array_api_tests/array_helpers.py | 10 ++++++----
 array_api_tests/test_linalg.py   | 14 ++++++++------
 2 files changed, 14 insertions(+), 10 deletions(-)

diff --git a/array_api_tests/array_helpers.py b/array_api_tests/array_helpers.py
index c34b87fb..9775ddac 100644
--- a/array_api_tests/array_helpers.py
+++ b/array_api_tests/array_helpers.py
@@ -205,7 +205,7 @@ 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.
 
@@ -213,11 +213,13 @@ def assert_exactly_equal(x, y):
     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)), f"The input arrays have different values ({x!r} != {y!r})"
+    assert all(exactly_equal(x, y)), f"The input arrays have different values ({x!r} != {y!r}){extra}"
 
 def assert_finite(x):
     """
diff --git a/array_api_tests/test_linalg.py b/array_api_tests/test_linalg.py
index 78b1be30..b027f96a 100644
--- a/array_api_tests/test_linalg.py
+++ b/array_api_tests/test_linalg.py
@@ -42,7 +42,8 @@
 
 pytestmark = pytest.mark.ci
 
-def assert_equal(x, y):
+def assert_equal(x, y, msg_extra=None):
+    extra = '' if not msg_extra else f' ({msg_extra})'
     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
@@ -51,10 +52,10 @@ def assert_equal(x, y):
 
         # 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})"
+        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}"
     else:
-        assert_exactly_equal(x, y)
+        assert_exactly_equal(x, y, msg_extra=msg_extra)
 
 def _test_stacks(f, *args, res=None, dims=2, true_val=None,
                  matrix_axes=(-2, -1),
@@ -93,9 +94,10 @@ def _test_stacks(f, *args, res=None, dims=2, true_val=None,
         res_stack = res[res_idx]
         x_stacks = [x[x_idx] for x, x_idx in zip(args, x_idxes)]
         decomp_res_stack = f(*x_stacks, **kw)
-        assert_equal(res_stack, decomp_res_stack)
+        msg_extra = f'{x_idxes = }, {res_idx = }'
+        assert_equal(res_stack, decomp_res_stack, msg_extra)
         if true_val:
-            assert_equal(decomp_res_stack, true_val(*x_stacks))
+            assert_equal(decomp_res_stack, true_val(*x_stacks), msg_extra)
 
 def _test_namedtuple(res, fields, func_name):
     """

From 1daba5d26d74507c0212cf581cda3caec522570a Mon Sep 17 00:00:00 2001
From: Aaron Meurer <asmeurer@gmail.com>
Date: Tue, 28 Mar 2023 19:46:41 -0600
Subject: [PATCH 30/54] Fix the true_value check for test_vecdot

---
 array_api_tests/test_linalg.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/array_api_tests/test_linalg.py b/array_api_tests/test_linalg.py
index b027f96a..d1424c61 100644
--- a/array_api_tests/test_linalg.py
+++ b/array_api_tests/test_linalg.py
@@ -97,7 +97,7 @@ def _test_stacks(f, *args, res=None, dims=2, true_val=None,
         msg_extra = f'{x_idxes = }, {res_idx = }'
         assert_equal(res_stack, decomp_res_stack, msg_extra)
         if true_val:
-            assert_equal(decomp_res_stack, true_val(*x_stacks), msg_extra)
+            assert_equal(decomp_res_stack, true_val(*x_stacks, **kw), msg_extra)
 
 def _test_namedtuple(res, fields, func_name):
     """
@@ -789,7 +789,7 @@ def test_vecdot(x1, x2, kw):
     ph.assert_shape("vecdot", res.shape, expected_shape)
 
     if x1.dtype in dh.int_dtypes:
-        def true_val(x, y, axix=-1):
+        def true_val(x, y, axis=-1):
             return xp.sum(x*y, dtype=res.dtype)
     else:
         true_val = None

From bbfe50fdbbde30f1a8a9c2d073e6438069732662 Mon Sep 17 00:00:00 2001
From: Aaron Meurer <asmeurer@gmail.com>
Date: Tue, 28 Mar 2023 19:47:52 -0600
Subject: [PATCH 31/54] Fix the test_diagonal true value check

---
 array_api_tests/test_linalg.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/array_api_tests/test_linalg.py b/array_api_tests/test_linalg.py
index d1424c61..1c7a1405 100644
--- a/array_api_tests/test_linalg.py
+++ b/array_api_tests/test_linalg.py
@@ -234,7 +234,7 @@ def test_diagonal(x, kw):
 
     assert res.shape == (*x.shape[:-2], diag_size), "diagonal() returned the wrong shape"
 
-    def true_diag(x_stack):
+    def true_diag(x_stack, offset=0):
         if offset >= 0:
             x_stack_diag = [x_stack[i, i + offset] for i in range(diag_size)]
         else:

From 64b0342dd25c087ea5151b2332e46cf280a1fa98 Mon Sep 17 00:00:00 2001
From: Aaron Meurer <asmeurer@gmail.com>
Date: Tue, 28 Mar 2023 19:49:50 -0600
Subject: [PATCH 32/54] Use a function instead of operation

Functions are easier to wrap in the compat layer, working around things like
type promotion differences.
---
 array_api_tests/test_linalg.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/array_api_tests/test_linalg.py b/array_api_tests/test_linalg.py
index 1c7a1405..39e5f670 100644
--- a/array_api_tests/test_linalg.py
+++ b/array_api_tests/test_linalg.py
@@ -790,7 +790,7 @@ def test_vecdot(x1, x2, kw):
 
     if x1.dtype in dh.int_dtypes:
         def true_val(x, y, axis=-1):
-            return xp.sum(x*y, dtype=res.dtype)
+            return xp.sum(xp.multiply(x, y), dtype=res.dtype)
     else:
         true_val = None
 

From 9cb58a1101adb1a97c653e3476adb36d1bbc3f44 Mon Sep 17 00:00:00 2001
From: Aaron Meurer <asmeurer@gmail.com>
Date: Tue, 18 Apr 2023 13:53:55 -0600
Subject: [PATCH 33/54] Add a comment

---
 array_api_tests/hypothesis_helpers.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/array_api_tests/hypothesis_helpers.py b/array_api_tests/hypothesis_helpers.py
index a5c06ed9..fee99c82 100644
--- a/array_api_tests/hypothesis_helpers.py
+++ b/array_api_tests/hypothesis_helpers.py
@@ -138,6 +138,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)

From c51216b6a505fd9bbb50dcf4752b7fb87e6eb153 Mon Sep 17 00:00:00 2001
From: Aaron Meurer <asmeurer@gmail.com>
Date: Fri, 2 Feb 2024 17:46:02 -0700
Subject: [PATCH 34/54] Remove flaky skips from linalg tests

---
 array_api_tests/test_linalg.py | 2 --
 1 file changed, 2 deletions(-)

diff --git a/array_api_tests/test_linalg.py b/array_api_tests/test_linalg.py
index 768d4b96..1f0b10f4 100644
--- a/array_api_tests/test_linalg.py
+++ b/array_api_tests/test_linalg.py
@@ -294,7 +294,6 @@ def test_eigvalsh(x):
 
     # TODO: Test that res actually corresponds to the eigenvalues of x
 
-@pytest.mark.skip(reason="flaky")
 @pytest.mark.xp_extension('linalg')
 @given(x=invertible_matrices())
 def test_inv(x):
@@ -307,7 +306,6 @@ def test_inv(x):
 
     # TODO: Test that the result is actually the inverse
 
-@pytest.mark.skip(reason="flaky")
 @given(
     *two_mutual_arrays(dh.real_dtypes)
 )

From cffd076dfbcc673e943113ad085a2921ef60cc99 Mon Sep 17 00:00:00 2001
From: Aaron Meurer <asmeurer@gmail.com>
Date: Fri, 2 Feb 2024 17:46:11 -0700
Subject: [PATCH 35/54] Fix some issues in linalg tests from recent merge

---
 array_api_tests/hypothesis_helpers.py |  1 +
 array_api_tests/test_linalg.py        | 18 +++++++++++-------
 2 files changed, 12 insertions(+), 7 deletions(-)

diff --git a/array_api_tests/hypothesis_helpers.py b/array_api_tests/hypothesis_helpers.py
index 563d6766..b963bd7f 100644
--- a/array_api_tests/hypothesis_helpers.py
+++ b/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
diff --git a/array_api_tests/test_linalg.py b/array_api_tests/test_linalg.py
index 1f0b10f4..c80c06c2 100644
--- a/array_api_tests/test_linalg.py
+++ b/array_api_tests/test_linalg.py
@@ -45,7 +45,7 @@
 
 def assert_equal(x, y, msg_extra=None):
     extra = '' if not msg_extra else f' ({msg_extra})'
-    if x.dtype in dh.float_dtypes:
+    if x.dtype in dh.all_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
@@ -701,7 +701,8 @@ def test_tensordot(x1, x2, kw):
     # TODO: vary shapes, vary contracted axes, test different axes arguments
     res = xp.tensordot(x1, x2, **kw)
 
-    ph.assert_dtype("tensordot", [x1.dtype, x2.dtype], res.dtype)
+    ph.assert_dtype("tensordot", in_dtype=[x1.dtype, x2.dtype],
+                    out_dtype=res.dtype)
 
     axes = _axes = kw.get('axes', 2)
 
@@ -785,9 +786,10 @@ def test_vecdot(x1, x2, kw):
 
     res = xp.vecdot(x1, x2, **kw)
 
-    ph.assert_dtype("vecdot", [x1.dtype, x2.dtype], res.dtype)
+    ph.assert_dtype("vecdot", in_dtype=[x1.dtype, x2.dtype],
+                    out_dtype=res.dtype)
     # TODO: assert shape and elements
-    ph.assert_shape("vecdot", res.shape, expected_shape)
+    ph.assert_shape("vecdot", out_shape=res.shape, expected=expected_shape)
 
     if x1.dtype in dh.int_dtypes:
         def true_val(x, y, axis=-1):
@@ -827,9 +829,11 @@ def test_vector_norm(x, data):
 
     _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)
+    ph.assert_keepdimable_shape('linalg.vector_norm', out_shape=res.shape,
+                                in_shape=x.shape, axes=_axes,
+                                keepdims=keepdims, kw=kw)
+    ph.assert_dtype('linalg.vector_norm', in_dtype=x.dtype,
+                    out_dtype=res.dtype)
 
     _kw = kw.copy()
     _kw.pop('axis', None)

From 3501116b36dfa257d8f2f8352416192d57220e35 Mon Sep 17 00:00:00 2001
From: Aaron Meurer <asmeurer@gmail.com>
Date: Fri, 2 Feb 2024 17:46:52 -0700
Subject: [PATCH 36/54] Fix vector_norm to not use our custom arrays strategy

---
 array_api_tests/test_linalg.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/array_api_tests/test_linalg.py b/array_api_tests/test_linalg.py
index c80c06c2..238abfdc 100644
--- a/array_api_tests/test_linalg.py
+++ b/array_api_tests/test_linalg.py
@@ -806,7 +806,7 @@ def true_val(x, y, axis=-1):
 
 @pytest.mark.xp_extension('linalg')
 @given(
-    x=xps.arrays(dtype=xps.floating_dtypes(), shape=shapes(min_side=1)),
+    x=arrays(dtype=xps.floating_dtypes(), shape=shapes(min_side=1)),
     data=data(),
 )
 def test_vector_norm(x, data):

From 5c1aa459833e54b34f9b421cb164bfd47b2c0660 Mon Sep 17 00:00:00 2001
From: Aaron Meurer <asmeurer@gmail.com>
Date: Fri, 2 Feb 2024 17:48:24 -0700
Subject: [PATCH 37/54] Update _test_stacks to use updated ndindex behavior

This requires https://github.com/Quansight-Labs/ndindex/pull/155 which is not
yet released.
---
 array_api_tests/test_linalg.py | 21 ++++++++++++++-------
 1 file changed, 14 insertions(+), 7 deletions(-)

diff --git a/array_api_tests/test_linalg.py b/array_api_tests/test_linalg.py
index 238abfdc..6d3ed9a3 100644
--- a/array_api_tests/test_linalg.py
+++ b/array_api_tests/test_linalg.py
@@ -60,6 +60,7 @@ def assert_equal(x, y, msg_extra=None):
 
 def _test_stacks(f, *args, res=None, dims=2, true_val=None,
                  matrix_axes=(-2, -1),
+                 res_axes=None,
                  assert_equal=assert_equal, **kw):
     """
     Test that f(*args, **kw) maps across stacks of matrices
@@ -84,7 +85,10 @@ def _test_stacks(f, *args, res=None, dims=2, true_val=None,
 
     # Assume the result is stacked along the last 'dims' axes of matrix_axes.
     # This holds for all the functions tested in this file
-    res_axes = matrix_axes[::-1][:dims]
+    if res_axes is None:
+        if not isinstance(matrix_axes, tuple) and all(isinstance(x, int) for x in matrix_axes):
+            raise ValueError("res_axes must be specified if matrix_axes is not a tuple of integers")
+        res_axes = matrix_axes[::-1][:dims]
 
     for (x_idxes, (res_idx,)) in zip(
             iter_indices(*shapes, skip_axes=matrix_axes),
@@ -330,10 +334,12 @@ def test_matmul(x1, x2):
         assert res.shape == ()
     elif len(x1.shape) == 1:
         assert res.shape == x2.shape[:-2] + x2.shape[-1:]
-        _test_stacks(_array_module.matmul, x1, x2, res=res, dims=1)
+        _test_stacks(_array_module.matmul, x1, x2, res=res, dims=1,
+                     matrix_axes=[(0,), (-2, -1)], res_axes=[-1])
     elif len(x2.shape) == 1:
         assert res.shape == x1.shape[:-1]
-        _test_stacks(_array_module.matmul, x1, x2, res=res, dims=1)
+        _test_stacks(_array_module.matmul, x1, x2, res=res, dims=1,
+                     matrix_axes=[(-2, -1), (0,)], res_axes=[-1])
     else:
         stack_shape = sh.broadcast_shapes(x1.shape[:-2], x2.shape[:-2])
         assert res.shape == stack_shape + (x1.shape[-2], x2.shape[-1])
@@ -546,10 +552,11 @@ def test_solve(x1, x2):
     # TODO: This requires an upstream fix to ndindex
     # (https://github.com/Quansight-Labs/ndindex/pull/131)
 
-    # if x2.ndim == 1:
-    #     _test_stacks(linalg.solve, x1, x2, res=res, dims=1)
-    # else:
-    #     _test_stacks(linalg.solve, x1, x2, res=res, dims=2)
+    if x2.ndim == 1:
+        _test_stacks(linalg.solve, x1, x2, res=res, dims=1,
+                     matrix_axes=[(-2, -1), (0,)], res_axes=[-1])
+    else:
+        _test_stacks(linalg.solve, x1, x2, res=res, dims=2)
 
 @pytest.mark.xp_extension('linalg')
 @given(

From 7a46e6b16d2b533182a7f49c03f489a1e39c2bca Mon Sep 17 00:00:00 2001
From: Aaron Meurer <asmeurer@gmail.com>
Date: Fri, 2 Feb 2024 17:53:15 -0700
Subject: [PATCH 38/54] Further limit the size of n in test_matrix_power

---
 array_api_tests/test_linalg.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/array_api_tests/test_linalg.py b/array_api_tests/test_linalg.py
index 6d3ed9a3..0c352a07 100644
--- a/array_api_tests/test_linalg.py
+++ b/array_api_tests/test_linalg.py
@@ -368,7 +368,7 @@ def test_matrix_norm(x, kw):
     _test_stacks(linalg.matrix_norm, x, **kw, dims=2 if keepdims else 0,
                  res=res)
 
-matrix_power_n = shared(integers(-1000, 1000), key='matrix_power n')
+matrix_power_n = shared(integers(-100, 100), key='matrix_power n')
 @pytest.mark.xp_extension('linalg')
 @given(
     # Generate any square matrix if n >= 0 but only invertible matrices if n < 0

From 6d154f29dd91b2dedf927526b5e7c39894dd2167 Mon Sep 17 00:00:00 2001
From: Aaron Meurer <asmeurer@gmail.com>
Date: Fri, 2 Feb 2024 17:53:27 -0700
Subject: [PATCH 39/54] Fix test_trace

---
 array_api_tests/test_linalg.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/array_api_tests/test_linalg.py b/array_api_tests/test_linalg.py
index 0c352a07..9343d23f 100644
--- a/array_api_tests/test_linalg.py
+++ b/array_api_tests/test_linalg.py
@@ -745,10 +745,9 @@ def test_trace(x, kw):
     # assert res.dtype == x.dtype, "trace() returned the wrong dtype"
 
     n, m = x.shape[-2:]
-    offset = kw.get('offset', 0)
     assert res.shape == x.shape[:-2], "trace() returned the wrong shape"
 
-    def true_trace(x_stack):
+    def true_trace(x_stack, offset=0):
         # Note: the spec does not specify that offset must be within the
         # bounds of the matrix. A large offset should just produce a size 0
         # diagonal in the last dimension (trace 0). See test_diagonal().

From 257aa138c053d96ac2e0b519695bb843e612b2d2 Mon Sep 17 00:00:00 2001
From: Aaron Meurer <asmeurer@gmail.com>
Date: Fri, 2 Feb 2024 17:57:37 -0700
Subject: [PATCH 40/54] Fix test_vecdot to only generate axis in [-min(x1.ndim,
 x2.ndim), -1]

See https://github.com/data-apis/array-api/pull/740
---
 array_api_tests/test_linalg.py | 10 ++++------
 1 file changed, 4 insertions(+), 6 deletions(-)

diff --git a/array_api_tests/test_linalg.py b/array_api_tests/test_linalg.py
index 9343d23f..4da74db6 100644
--- a/array_api_tests/test_linalg.py
+++ b/array_api_tests/test_linalg.py
@@ -769,17 +769,15 @@ def true_trace(x_stack, offset=0):
 
 @given(
     *two_mutual_arrays(dh.numeric_dtypes, mutually_broadcastable_shapes(2, min_dims=1)),
-    kwargs(axis=integers()),
+    data(),
 )
-def test_vecdot(x1, x2, kw):
+def test_vecdot(x1, x2, data):
     # TODO: vary shapes, test different axis arguments
     broadcasted_shape = sh.broadcast_shapes(x1.shape, x2.shape)
+    min_ndim = min(x1.ndim, x2.ndim)
     ndim = len(broadcasted_shape)
+    kw = data.draw(kwargs(axis=integers(-min_ndim, -1)))
     axis = kw.get('axis', -1)
-    if not (-ndim <= axis < ndim):
-        ph.raises(Exception, lambda: xp.vecdot(x1, x2, **kw),
-                  f"vecdot did not raise an exception for invalid axis ({ndim=}, {kw=})")
-        return
     x1_shape = (1,)*(ndim - x1.ndim) + tuple(x1.shape)
     x2_shape = (1,)*(ndim - x2.ndim) + tuple(x2.shape)
     if x1_shape[axis] != x2_shape[axis]:

From afc8a258c3e2a73c64d4e378a7c9c9d3b90ca2b6 Mon Sep 17 00:00:00 2001
From: Aaron Meurer <asmeurer@gmail.com>
Date: Fri, 2 Feb 2024 18:07:55 -0700
Subject: [PATCH 41/54] Update test_cross to test broadcastable shapes

This also updates it to only test axes from [min(x1.ndim, x2.ndim), -1], as
per https://github.com/data-apis/array-api/pull/740
---
 array_api_tests/test_linalg.py | 34 +++++++++++++++++++++-------------
 1 file changed, 21 insertions(+), 13 deletions(-)

diff --git a/array_api_tests/test_linalg.py b/array_api_tests/test_linalg.py
index 4da74db6..d8ca557a 100644
--- a/array_api_tests/test_linalg.py
+++ b/array_api_tests/test_linalg.py
@@ -148,23 +148,29 @@ def cross_args(draw, dtype_objects=dh.real_dtypes):
     in the drawn axis.
 
     """
-    shape = list(draw(shapes()))
-    size = len(shape)
-    assume(size > 0)
+    shape1, shape2 = draw(two_mutually_broadcastable_shapes)
+    min_ndim = min(len(shape1), len(shape2))
+    assume(min_ndim > 0)
 
-    kw = draw(kwargs(axis=integers(-size, size-1)))
+    kw = draw(kwargs(axis=integers(-min_ndim, -1)))
     axis = kw.get('axis', -1)
-    shape[axis] = 3
-    shape = tuple(shape)
+    if draw(booleans()):
+        # Sometimes allow invalid inputs to test it errors
+        shape1 = list(shape1)
+        shape1[axis] = 3
+        shape1 = tuple(shape1)
+        shape2 = list(shape2)
+        shape2[axis] = 3
+        shape2 = tuple(shape2)
 
     mutual_dtypes = shared(mutually_promotable_dtypes(dtypes=dtype_objects))
     arrays1 = arrays(
         dtype=mutual_dtypes.map(lambda pair: pair[0]),
-        shape=shape,
+        shape=shape1,
     )
     arrays2 = arrays(
         dtype=mutual_dtypes.map(lambda pair: pair[1]),
-        shape=shape,
+        shape=shape2,
     )
     return draw(arrays1), draw(arrays2), kw
 
@@ -176,15 +182,17 @@ def test_cross(x1_x2_kw):
     x1, x2, kw = x1_x2_kw
 
     axis = kw.get('axis', -1)
-    err = "test_cross produced invalid input. This indicates a bug in the test suite."
-    assert x1.shape == x2.shape, err
-    shape = x1.shape
-    assert x1.shape[axis] == x2.shape[axis] == 3, err
+    if not (x1.shape[axis] == x2.shape[axis] == 3):
+        ph.raises(Exception, lambda: xp.cross(x1, x2, **kw),
+                  "cross did not raise an exception for invalid shapes")
+        return
 
     res = linalg.cross(x1, x2, **kw)
 
+    broadcasted_shape = sh.broadcast_shapes(x1.shape, x2.shape)
+
     assert res.dtype == dh.result_type(x1.dtype, x2.dtype), "cross() did not return the correct dtype"
-    assert res.shape == shape, "cross() did not return the correct shape"
+    assert res.shape == broadcasted_shape, "cross() did not return the correct shape"
 
     def exact_cross(a, b):
         assert a.shape == b.shape == (3,), "Invalid cross() stack shapes. This indicates a bug in the test suite."

From 3cb991208251e47eca24010d997ea4942674291b Mon Sep 17 00:00:00 2001
From: Aaron Meurer <asmeurer@gmail.com>
Date: Fri, 2 Feb 2024 18:09:20 -0700
Subject: [PATCH 42/54] Fix test_cross to use assert_dtype and assert_shape
 helpers

---
 array_api_tests/test_linalg.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/array_api_tests/test_linalg.py b/array_api_tests/test_linalg.py
index d8ca557a..7982052a 100644
--- a/array_api_tests/test_linalg.py
+++ b/array_api_tests/test_linalg.py
@@ -191,8 +191,9 @@ def test_cross(x1_x2_kw):
 
     broadcasted_shape = sh.broadcast_shapes(x1.shape, x2.shape)
 
-    assert res.dtype == dh.result_type(x1.dtype, x2.dtype), "cross() did not return the correct dtype"
-    assert res.shape == broadcasted_shape, "cross() did not return the correct shape"
+    ph.assert_dtype("cross", in_dtype=[x1.dtype, x2.dtype],
+                    out_dtype=res.dtype)
+    ph.assert_shape("cross", out_shape=res.shape, expected=broadcasted_shape)
 
     def exact_cross(a, b):
         assert a.shape == b.shape == (3,), "Invalid cross() stack shapes. This indicates a bug in the test suite."
@@ -800,7 +801,6 @@ def test_vecdot(x1, x2, data):
 
     ph.assert_dtype("vecdot", in_dtype=[x1.dtype, x2.dtype],
                     out_dtype=res.dtype)
-    # TODO: assert shape and elements
     ph.assert_shape("vecdot", out_shape=res.shape, expected=expected_shape)
 
     if x1.dtype in dh.int_dtypes:

From 012ca19a35059442446449524150267ff59fbb8c Mon Sep 17 00:00:00 2001
From: Aaron Meurer <asmeurer@gmail.com>
Date: Fri, 2 Feb 2024 18:12:00 -0700
Subject: [PATCH 43/54] Remove some completed TODO comments

---
 array_api_tests/test_linalg.py | 8 +-------
 1 file changed, 1 insertion(+), 7 deletions(-)

diff --git a/array_api_tests/test_linalg.py b/array_api_tests/test_linalg.py
index 7982052a..5bc19265 100644
--- a/array_api_tests/test_linalg.py
+++ b/array_api_tests/test_linalg.py
@@ -12,7 +12,7 @@
 required, but we don't yet have a clean way to disable only those tests (see https://github.com/data-apis/array-api-tests/issues/25).
 
 """
-# TODO: test with complex dtypes where appropiate
+# TODO: test with complex dtypes where appropriate
 
 import pytest
 from hypothesis import assume, given
@@ -558,9 +558,6 @@ def _x2_shapes(draw):
 def test_solve(x1, x2):
     res = linalg.solve(x1, x2)
 
-    # TODO: This requires an upstream fix to ndindex
-    # (https://github.com/Quansight-Labs/ndindex/pull/131)
-
     if x2.ndim == 1:
         _test_stacks(linalg.solve, x1, x2, res=res, dims=1,
                      matrix_axes=[(-2, -1), (0,)], res_axes=[-1])
@@ -714,7 +711,6 @@ def _test_tensordot_stacks(x1, x2, kw, res):
     tensordot_kw,
 )
 def test_tensordot(x1, x2, kw):
-    # TODO: vary shapes, vary contracted axes, test different axes arguments
     res = xp.tensordot(x1, x2, **kw)
 
     ph.assert_dtype("tensordot", in_dtype=[x1.dtype, x2.dtype],
@@ -734,7 +730,6 @@ def test_tensordot(x1, x2, kw):
     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')
@@ -781,7 +776,6 @@ def true_trace(x_stack, offset=0):
     data(),
 )
 def test_vecdot(x1, x2, data):
-    # TODO: vary shapes, test different axis arguments
     broadcasted_shape = sh.broadcast_shapes(x1.shape, x2.shape)
     min_ndim = min(x1.ndim, x2.ndim)
     ndim = len(broadcasted_shape)

From 5ceb81d87194f528dab04daedcdeb09693a9065b Mon Sep 17 00:00:00 2001
From: Aaron Meurer <asmeurer@gmail.com>
Date: Fri, 2 Feb 2024 18:40:02 -0700
Subject: [PATCH 44/54] Update linalg tests to test complex dtypes

Also diagonal and matrix_transpose now test against all dtypes, since they
have no dtype restrictions in the spec.
---
 array_api_tests/dtype_helpers.py | 10 ++++++++
 array_api_tests/test_linalg.py   | 43 ++++++++++++++++++--------------
 2 files changed, 34 insertions(+), 19 deletions(-)

diff --git a/array_api_tests/dtype_helpers.py b/array_api_tests/dtype_helpers.py
index 3052d54f..9d26cd04 100644
--- a/array_api_tests/dtype_helpers.py
+++ b/array_api_tests/dtype_helpers.py
@@ -231,6 +231,16 @@ 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")
 
 if not hasattr(xp, "asarray"):
     default_int = xp.int32
diff --git a/array_api_tests/test_linalg.py b/array_api_tests/test_linalg.py
index 5bc19265..3a274c0f 100644
--- a/array_api_tests/test_linalg.py
+++ b/array_api_tests/test_linalg.py
@@ -12,8 +12,6 @@
 required, but we don't yet have a clean way to disable only those tests (see https://github.com/data-apis/array-api-tests/issues/25).
 
 """
-# TODO: test with complex dtypes where appropriate
-
 import pytest
 from hypothesis import assume, given
 from hypothesis.strategies import (booleans, composite, tuples, floats,
@@ -24,8 +22,9 @@
 import itertools
 
 from .array_helpers import assert_exactly_equal, asarray
-from .hypothesis_helpers import (arrays, xps, shapes, kwargs, matrix_shapes,
-                                 square_matrix_shapes, symmetric_matrices,
+from .hypothesis_helpers import (arrays, all_floating_dtypes, xps, 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,
@@ -210,7 +209,7 @@ def exact_cross(a, b):
 
 @pytest.mark.xp_extension('linalg')
 @given(
-    x=arrays(dtype=xps.floating_dtypes(), shape=square_matrix_shapes),
+    x=arrays(dtype=all_floating_dtypes(), shape=square_matrix_shapes),
 )
 def test_det(x):
     res = linalg.det(x)
@@ -224,7 +223,7 @@ def test_det(x):
 
 @pytest.mark.xp_extension('linalg')
 @given(
-    x=arrays(dtype=xps.real_dtypes(), shape=matrix_shapes()),
+    x=arrays(dtype=xps.scalar_dtypes(), shape=matrix_shapes()),
     # offset may produce an overflow if it is too large. Supporting offsets
     # that are way larger than the array shape isn't very important.
     kw=kwargs(offset=integers(-MAX_ARRAY_SIZE, MAX_ARRAY_SIZE))
@@ -382,7 +381,7 @@ def test_matrix_norm(x, kw):
 @given(
     # Generate any square matrix if n >= 0 but only invertible matrices if n < 0
     x=matrix_power_n.flatmap(lambda n: invertible_matrices() if n < 0 else
-                             arrays(dtype=xps.floating_dtypes(),
+                             arrays(dtype=all_floating_dtypes(),
                                         shape=square_matrix_shapes)),
     n=matrix_power_n,
 )
@@ -409,7 +408,7 @@ def test_matrix_rank(x, kw):
     linalg.matrix_rank(x, **kw)
 
 @given(
-    x=arrays(dtype=xps.real_dtypes(), shape=matrix_shapes()),
+    x=arrays(dtype=xps.scalar_dtypes(), shape=matrix_shapes()),
 )
 def test_matrix_transpose(x):
     res = _array_module.matrix_transpose(x)
@@ -459,7 +458,7 @@ def test_pinv(x, kw):
 
 @pytest.mark.xp_extension('linalg')
 @given(
-    x=arrays(dtype=xps.floating_dtypes(), shape=matrix_shapes()),
+    x=arrays(dtype=all_floating_dtypes(), shape=matrix_shapes()),
     kw=kwargs(mode=sampled_from(['reduced', 'complete']))
 )
 def test_qr(x, kw):
@@ -495,7 +494,7 @@ def test_qr(x, kw):
 
 @pytest.mark.xp_extension('linalg')
 @given(
-    x=arrays(dtype=xps.floating_dtypes(), shape=square_matrix_shapes),
+    x=arrays(dtype=all_floating_dtypes(), shape=square_matrix_shapes),
 )
 def test_slogdet(x):
     res = linalg.slogdet(x)
@@ -504,11 +503,16 @@ def test_slogdet(x):
 
     sign, logabsdet = res
 
-    assert sign.dtype == x.dtype, "slogdet().sign did not return the correct dtype"
-    assert sign.shape == x.shape[:-2], "slogdet().sign did not return the correct shape"
-    assert logabsdet.dtype == x.dtype, "slogdet().logabsdet did not return the correct dtype"
-    assert logabsdet.shape == x.shape[:-2], "slogdet().logabsdet did not return the correct shape"
-
+    ph.assert_dtype("slogdet", in_dtype=x.dtype, out_dtype=sign.dtype,
+                    expected=x.dtype, repr_name="sign.dtype")
+    ph.assert_shape("slogdet", out_shape=sign.shape, expected=x.shape[:-2],
+                    repr_name="sign.shape")
+    expected_dtype = dh.as_real_dtype(x.dtype)
+    ph.assert_dtype("slogdet", in_dtype=x.dtype, out_dtype=logabsdet.dtype,
+                    expected=expected_dtype, repr_name="logabsdet.dtype")
+    ph.assert_shape("slogdet", out_shape=logabsdet.shape,
+                    expected=x.shape[:-2],
+                    repr_name="logabsdet.shape")
 
     _test_stacks(lambda x: linalg.slogdet(x).sign, x,
                  res=sign, dims=0)
@@ -550,7 +554,7 @@ def _x2_shapes(draw):
         return draw(stack_shapes)[1] + draw(x1).shape[-1:] + (end,)
 
     x2_shapes = one_of(x1.map(lambda x: (x.shape[-1],)), _x2_shapes())
-    x2 = arrays(dtype=xps.floating_dtypes(), shape=x2_shapes)
+    x2 = arrays(dtype=all_floating_dtypes(), shape=x2_shapes)
     return x1, x2
 
 @pytest.mark.xp_extension('linalg')
@@ -734,7 +738,7 @@ def test_tensordot(x1, x2, kw):
 
 @pytest.mark.xp_extension('linalg')
 @given(
-    x=arrays(dtype=xps.real_dtypes(), shape=matrix_shapes()),
+    x=arrays(dtype=xps.numeric_dtypes(), shape=matrix_shapes()),
     # offset may produce an overflow if it is too large. Supporting offsets
     # that are way larger than the array shape isn't very important.
     kw=kwargs(offset=integers(-MAX_ARRAY_SIZE, MAX_ARRAY_SIZE))
@@ -812,7 +816,7 @@ def true_val(x, y, axis=-1):
 
 @pytest.mark.xp_extension('linalg')
 @given(
-    x=arrays(dtype=xps.floating_dtypes(), shape=shapes(min_side=1)),
+    x=arrays(dtype=all_floating_dtypes(), shape=shapes(min_side=1)),
     data=data(),
 )
 def test_vector_norm(x, data):
@@ -838,8 +842,9 @@ def test_vector_norm(x, data):
     ph.assert_keepdimable_shape('linalg.vector_norm', out_shape=res.shape,
                                 in_shape=x.shape, axes=_axes,
                                 keepdims=keepdims, kw=kw)
+    expected_dtype = dh.as_real_dtype(x.dtype)
     ph.assert_dtype('linalg.vector_norm', in_dtype=x.dtype,
-                    out_dtype=res.dtype)
+                    out_dtype=res.dtype, expected=expected_dtype)
 
     _kw = kw.copy()
     _kw.pop('axis', None)

From a4d419ff6da386c9dcaef16096dcce7c27082d75 Mon Sep 17 00:00:00 2001
From: Aaron Meurer <asmeurer@gmail.com>
Date: Fri, 2 Feb 2024 18:58:35 -0700
Subject: [PATCH 45/54] Update linalg tests to use assert_dtype and
 assert_shape helpers

---
 array_api_tests/test_linalg.py | 150 ++++++++++++++++++++++-----------
 1 file changed, 100 insertions(+), 50 deletions(-)

diff --git a/array_api_tests/test_linalg.py b/array_api_tests/test_linalg.py
index 3a274c0f..096abe44 100644
--- a/array_api_tests/test_linalg.py
+++ b/array_api_tests/test_linalg.py
@@ -126,8 +126,9 @@ def _test_namedtuple(res, fields, func_name):
 def test_cholesky(x, kw):
     res = linalg.cholesky(x, **kw)
 
-    assert res.shape == x.shape, "cholesky() did not return the correct shape"
-    assert res.dtype == x.dtype, "cholesky() did not return the correct dtype"
+    ph.assert_dtype("cholesky", in_dtype=x.dtype, out_dtype=res.dtype)
+    ph.assert_result_shape("cholesky", in_shapes=[x.shape],
+                           out_shape=res.shape, expected=x.shape)
 
     _test_stacks(linalg.cholesky, x, **kw, res=res)
 
@@ -192,7 +193,7 @@ def test_cross(x1_x2_kw):
 
     ph.assert_dtype("cross", in_dtype=[x1.dtype, x2.dtype],
                     out_dtype=res.dtype)
-    ph.assert_shape("cross", out_shape=res.shape, expected=broadcasted_shape)
+    ph.assert_result_shape("cross", in_shapes=[x1.shape, x2.shape], out_shape=res.shape, expected=broadcasted_shape)
 
     def exact_cross(a, b):
         assert a.shape == b.shape == (3,), "Invalid cross() stack shapes. This indicates a bug in the test suite."
@@ -214,8 +215,9 @@ def exact_cross(a, b):
 def test_det(x):
     res = linalg.det(x)
 
-    assert res.dtype == x.dtype, "det() did not return the correct dtype"
-    assert res.shape == x.shape[:-2], "det() did not return the correct shape"
+    ph.assert_dtype("det", in_dtype=x.dtype, out_dtype=res.dtype)
+    ph.assert_result_shape("det", in_shapes=[x.shape], out_shape=res.shape,
+                           expected=x.shape[:-2])
 
     _test_stacks(linalg.det, x, res=res, dims=0)
 
@@ -231,7 +233,7 @@ def test_det(x):
 def test_diagonal(x, kw):
     res = linalg.diagonal(x, **kw)
 
-    assert res.dtype == x.dtype, "diagonal() returned the wrong dtype"
+    ph.assert_dtype("diagonal", in_dtype=x.dtype, out_dtype=res.dtype)
 
     n, m = x.shape[-2:]
     offset = kw.get('offset', 0)
@@ -245,7 +247,9 @@ def test_diagonal(x, kw):
     else:
         diag_size = min(n, m, max(m - offset, 0))
 
-    assert res.shape == (*x.shape[:-2], diag_size), "diagonal() returned the wrong shape"
+    expected_shape = (*x.shape[:-2], diag_size)
+    ph.assert_result_shape("diagonal", in_shapes=[x.shape],
+                           out_shape=res.shape, expected=expected_shape)
 
     def true_diag(x_stack, offset=0):
         if offset >= 0:
@@ -266,11 +270,18 @@ def test_eigh(x):
     eigenvalues = res.eigenvalues
     eigenvectors = res.eigenvectors
 
-    assert eigenvalues.dtype == x.dtype, "eigh().eigenvalues did not return the correct dtype"
-    assert eigenvalues.shape == x.shape[:-1], "eigh().eigenvalues did not return the correct shape"
+    ph.assert_dtype("eigh", in_dtype=x.dtype, out_dtype=eigenvalues.dtype,
+                    expected=x.dtype, repr_name="eigenvalues.dtype")
+    ph.assert_result_shape("eigh", in_shapes=[x.shape],
+                           out_shape=eigenvalues.shape,
+                           expected=x.shape[:-1],
+                           repr_name="eigenvalues.shape")
 
-    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"
+    ph.assert_dtype("eigh", in_dtype=x.dtype, out_dtype=eigenvectors.dtype,
+                    expected=x.dtype, repr_name="eigenvectors.dtype")
+    ph.assert_result_shape("eigh", in_shapes=[x.shape],
+                           out_shape=eigenvectors.shape, expected=x.shape,
+                           repr_name="eigenvectors.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
@@ -292,8 +303,9 @@ def test_eigh(x):
 def test_eigvalsh(x):
     res = linalg.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"
+    ph.assert_dtype("eigvalsh", in_dtype=x.dtype, out_dtype=res.dtype)
+    ph.assert_result_shape("eigvalsh", in_shapes=[x.shape],
+                           out_shape=res.shape, expected=x.shape[:-1])
 
     # 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
@@ -311,8 +323,9 @@ def test_eigvalsh(x):
 def test_inv(x):
     res = linalg.inv(x)
 
-    assert res.shape == x.shape, "inv() did not return the correct shape"
-    assert res.dtype == x.dtype, "inv() did not return the correct dtype"
+    ph.assert_dtype("inv", in_dtype=x.dtype, out_dtype=res.dtype)
+    ph.assert_result_shape("inv", in_shapes=[x.shape], out_shape=res.shape,
+                           expected=x.shape)
 
     _test_stacks(linalg.inv, x, res=res)
 
@@ -339,18 +352,24 @@ def test_matmul(x1, x2):
     ph.assert_dtype("matmul", in_dtype=[x1.dtype, x2.dtype], out_dtype=res.dtype)
 
     if len(x1.shape) == len(x2.shape) == 1:
-        assert res.shape == ()
+        ph.assert_result_shape("matmul", in_shapes=[x1.shape, x2.shape],
+                               out_shape=res.shape, expected=())
     elif len(x1.shape) == 1:
-        assert res.shape == x2.shape[:-2] + x2.shape[-1:]
+        ph.assert_result_shape("matmul", in_shapes=[x1.shape, x2.shape],
+                               out_shape=res.shape,
+                               expected=x2.shape[:-2] + x2.shape[-1:])
         _test_stacks(_array_module.matmul, x1, x2, res=res, dims=1,
                      matrix_axes=[(0,), (-2, -1)], res_axes=[-1])
     elif len(x2.shape) == 1:
-        assert res.shape == x1.shape[:-1]
+        ph.assert_result_shape("matmul", in_shapes=[x1.shape, x2.shape],
+                               out_shape=res.shape, expected=x1.shape[:-1])
         _test_stacks(_array_module.matmul, x1, x2, res=res, dims=1,
                      matrix_axes=[(-2, -1), (0,)], res_axes=[-1])
     else:
         stack_shape = sh.broadcast_shapes(x1.shape[:-2], x2.shape[:-2])
-        assert res.shape == stack_shape + (x1.shape[-2], x2.shape[-1])
+        ph.assert_result_shape("matmul", in_shapes=[x1.shape, x2.shape],
+                               out_shape=res.shape,
+                               expected=stack_shape + (x1.shape[-2], x2.shape[-1]))
         _test_stacks(_array_module.matmul, x1, x2, res=res)
 
 @pytest.mark.xp_extension('linalg')
@@ -370,8 +389,9 @@ def test_matrix_norm(x, kw):
         expected_shape = x.shape[:-2] + (1, 1)
     else:
         expected_shape = x.shape[:-2]
-    assert res.shape == expected_shape, f"matrix_norm({keepdims=}) did not return the correct shape"
-    assert res.dtype == x.dtype, "matrix_norm() did not return the correct dtype"
+    ph.assert_dtype("matrix_norm", in_dtype=x.dtype, out_dtype=res.dtype)
+    ph.assert_result_shape("matrix_norm", in_shapes=[x.shape],
+                           out_shape=res.shape, expected=expected_shape)
 
     _test_stacks(linalg.matrix_norm, x, **kw, dims=2 if keepdims else 0,
                  res=res)
@@ -388,8 +408,9 @@ def test_matrix_norm(x, kw):
 def test_matrix_power(x, n):
     res = linalg.matrix_power(x, n)
 
-    assert res.shape == x.shape, "matrix_power() did not return the correct shape"
-    assert res.dtype == x.dtype, "matrix_power() did not return the correct dtype"
+    ph.assert_dtype("matrix_power", in_dtype=x.dtype, out_dtype=res.dtype)
+    ph.assert_result_shape("matrix_power", in_shapes=[x.shape],
+                           out_shape=res.shape, expected=x.shape)
 
     if n == 0:
         true_val = lambda x: _array_module.eye(x.shape[0], dtype=x.dtype)
@@ -419,8 +440,9 @@ def test_matrix_transpose(x):
     shape = list(x.shape)
     shape[-1], shape[-2] = shape[-2], shape[-1]
     shape = tuple(shape)
-    assert res.shape == shape, "matrix_transpose() did not return the correct shape"
-    assert res.dtype == x.dtype, "matrix_transpose() did not return the correct dtype"
+    ph.assert_dtype("matrix_transpose", in_dtype=x.dtype, out_dtype=res.dtype)
+    ph.assert_result_shape("matrix_transpose", in_shapes=[x.shape],
+                           out_shape=res.shape, expected=shape)
 
     _test_stacks(_array_module.matrix_transpose, x, res=res, true_val=true_val)
 
@@ -435,8 +457,9 @@ def test_outer(x1, x2):
     res = linalg.outer(x1, x2)
 
     shape = (x1.shape[0], x2.shape[0])
-    assert res.shape == shape, "outer() did not return the correct shape"
-    assert res.dtype == dh.result_type(x1.dtype, x2.dtype), "outer() did not return the correct dtype"
+    ph.assert_dtype("outer", in_dtype=[x1.dtype, x2.dtype], out_dtype=res.dtype)
+    ph.assert_result_shape("outer", in_shapes=[x1.shape, x2.shape],
+                           out_shape=res.shape, expected=shape)
 
     if 0 in shape:
         true_res = _array_module.empty(shape, dtype=res.dtype)
@@ -472,17 +495,23 @@ def test_qr(x, kw):
     Q = res.Q
     R = res.R
 
-    assert Q.dtype == x.dtype, "qr().Q did not return the correct dtype"
+    ph.assert_dtype("qr", in_dtype=x.dtype, out_dtype=Q.dtype,
+                    expected=x.dtype, repr_name="Q.dtype")
     if mode == 'complete':
-        assert Q.shape == x.shape[:-2] + (M, M), "qr().Q did not return the correct shape"
+        expected_Q_shape = x.shape[:-2] + (M, M)
     else:
-        assert Q.shape == x.shape[:-2] + (M, K), "qr().Q did not return the correct shape"
+        expected_Q_shape = x.shape[:-2] + (M, K)
+    ph.assert_result_shape("qr", in_shapes=[x.shape], out_shape=Q.shape,
+                           expected=expected_Q_shape, repr_name="Q.shape")
 
-    assert R.dtype == x.dtype, "qr().R did not return the correct dtype"
+    ph.assert_dtype("qr", in_dtype=x.dtype, out_dtype=R.dtype,
+                    expected=x.dtype, repr_name="R.dtype")
     if mode == 'complete':
-        assert R.shape == x.shape[:-2] + (M, N), "qr().R did not return the correct shape"
+        expected_R_shape = x.shape[:-2] + (M, N)
     else:
-        assert R.shape == x.shape[:-2] + (K, N), "qr().R did not return the correct shape"
+        expected_R_shape = x.shape[:-2] + (K, N)
+    ph.assert_result_shape("qr", in_shapes=[x.shape], out_shape=R.shape,
+                           expected=expected_R_shape, repr_name="R.shape")
 
     _test_stacks(lambda x: linalg.qr(x, **kw).Q, x, res=Q)
     _test_stacks(lambda x: linalg.qr(x, **kw).R, x, res=R)
@@ -505,14 +534,17 @@ def test_slogdet(x):
 
     ph.assert_dtype("slogdet", in_dtype=x.dtype, out_dtype=sign.dtype,
                     expected=x.dtype, repr_name="sign.dtype")
-    ph.assert_shape("slogdet", out_shape=sign.shape, expected=x.shape[:-2],
-                    repr_name="sign.shape")
+    ph.assert_result_shape("slogdet", in_shapes=[x.shape],
+                           out_shape=sign.shape,
+                           expected=x.shape[:-2],
+                           repr_name="sign.shape")
     expected_dtype = dh.as_real_dtype(x.dtype)
     ph.assert_dtype("slogdet", in_dtype=x.dtype, out_dtype=logabsdet.dtype,
                     expected=expected_dtype, repr_name="logabsdet.dtype")
-    ph.assert_shape("slogdet", out_shape=logabsdet.shape,
-                    expected=x.shape[:-2],
-                    repr_name="logabsdet.shape")
+    ph.assert_result_shape("slogdet", in_shapes=[x.shape],
+                           out_shape=logabsdet.shape,
+                           expected=x.shape[:-2],
+                           repr_name="logabsdet.shape")
 
     _test_stacks(lambda x: linalg.slogdet(x).sign, x,
                  res=sign, dims=0)
@@ -584,17 +616,31 @@ def test_svd(x, kw):
 
     U, S, Vh = res
 
-    assert U.dtype == x.dtype, "svd().U did not return the correct dtype"
-    assert S.dtype == x.dtype, "svd().S did not return the correct dtype"
-    assert Vh.dtype == x.dtype, "svd().Vh did not return the correct dtype"
+    ph.assert_dtype("svd", in_dtype=x.dtype, out_dtype=U.dtype,
+                    expected=x.dtype, repr_name="U.dtype")
+    ph.assert_dtype("svd", in_dtype=x.dtype, out_dtype=S.dtype,
+                    expected=x.dtype, repr_name="S.dtype")
+    ph.assert_dtype("svd", in_dtype=x.dtype, out_dtype=Vh.dtype,
+                    expected=x.dtype, repr_name="Vh.dtype")
 
     if full_matrices:
-        assert U.shape == (*stack, M, M), "svd().U did not return the correct shape"
-        assert Vh.shape == (*stack, N, N), "svd().Vh did not return the correct shape"
+        expected_U_shape = (*stack, M, M)
+        expected_Vh_shape = (*stack, N, N)
     else:
-        assert U.shape == (*stack, M, K), "svd(full_matrices=False).U did not return the correct shape"
-        assert Vh.shape == (*stack, K, N), "svd(full_matrices=False).Vh did not return the correct shape"
-    assert S.shape == (*stack, K), "svd().S did not return the correct shape"
+        expected_U_shape = (*stack, M, K)
+        expected_Vh_shape = (*stack, K, N)
+    ph.assert_result_shape("svd", in_shapes=[x.shape],
+                           out_shape=U.shape,
+                           expected=expected_U_shape,
+                           repr_name="U.shape")
+    ph.assert_result_shape("svd", in_shapes=[x.shape],
+                            out_shape=Vh.shape,
+                            expected=expected_Vh_shape,
+                            repr_name="Vh.shape")
+    ph.assert_result_shape("svd", in_shapes=[x.shape],
+                            out_shape=S.shape,
+                            expected=(*stack, K),
+                            repr_name="S.shape")
 
     # The values of s must be sorted from largest to smallest
     if K >= 1:
@@ -614,8 +660,11 @@ def test_svdvals(x):
     *stack, M, N = x.shape
     K = min(M, N)
 
-    assert res.dtype == x.dtype, "svdvals() did not return the correct dtype"
-    assert res.shape == (*stack, K), "svdvals() did not return the correct shape"
+    ph.assert_dtype("svdvals", in_dtype=x.dtype, out_dtype=res.dtype,
+                    expected=x.dtype)
+    ph.assert_result_shape("svdvals", in_shapes=[x.shape],
+                           out_shape=res.shape,
+                           expected=(*stack, K))
 
     # SVD values must be sorted from largest to smallest
     assert _array_module.all(res[..., :-1] >= res[..., 1:]), "svdvals() values are not sorted from largest to smallest"
@@ -753,7 +802,7 @@ def test_trace(x, kw):
     # assert res.dtype == x.dtype, "trace() returned the wrong dtype"
 
     n, m = x.shape[-2:]
-    assert res.shape == x.shape[:-2], "trace() returned the wrong shape"
+    ph.assert_result_shape('trace', x.shape, res.shape, expected=x.shape[:-2])
 
     def true_trace(x_stack, offset=0):
         # Note: the spec does not specify that offset must be within the
@@ -799,7 +848,8 @@ def test_vecdot(x1, x2, data):
 
     ph.assert_dtype("vecdot", in_dtype=[x1.dtype, x2.dtype],
                     out_dtype=res.dtype)
-    ph.assert_shape("vecdot", out_shape=res.shape, expected=expected_shape)
+    ph.assert_result_shape("vecdot", in_shapes=[x1.shape, x2.shape],
+                           out_shape=res.shape, expected=expected_shape)
 
     if x1.dtype in dh.int_dtypes:
         def true_val(x, y, axis=-1):

From 6f9db94db029a851b567f7a53fc92cb9efc04868 Mon Sep 17 00:00:00 2001
From: Aaron Meurer <asmeurer@gmail.com>
Date: Fri, 2 Feb 2024 19:07:27 -0700
Subject: [PATCH 46/54] Factor out dtype logic from test_sum() and test_prod()
 and apply it to test_trace()

---
 array_api_tests/dtype_helpers.py              | 41 +++++++++++
 array_api_tests/test_linalg.py                | 15 ++--
 array_api_tests/test_statistical_functions.py | 70 ++-----------------
 3 files changed, 57 insertions(+), 69 deletions(-)

diff --git a/array_api_tests/dtype_helpers.py b/array_api_tests/dtype_helpers.py
index 9d26cd04..59edfe86 100644
--- a/array_api_tests/dtype_helpers.py
+++ b/array_api_tests/dtype_helpers.py
@@ -242,6 +242,47 @@ def as_real_dtype(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
     default_float = xp.float32
diff --git a/array_api_tests/test_linalg.py b/array_api_tests/test_linalg.py
index 096abe44..62862bba 100644
--- a/array_api_tests/test_linalg.py
+++ b/array_api_tests/test_linalg.py
@@ -795,11 +795,16 @@ def test_tensordot(x1, x2, kw):
 def test_trace(x, kw):
     res = linalg.trace(x, **kw)
 
-    # TODO: trace() should promote in some cases. See
-    # https://github.com/data-apis/array-api/issues/202. See also the dtype
-    # argument to sum() below.
-
-    # assert res.dtype == x.dtype, "trace() returned the wrong dtype"
+    dtype = kw.get("dtype", None)
+    expected_dtype = dh.accumulation_result_dtype(x.dtype, dtype)
+    if expected_dtype is None:
+        # If a default uint cannot exist (i.e. in PyTorch which doesn't support
+        # uint32 or uint64), we skip testing the output dtype.
+        # See https://github.com/data-apis/array-api-tests/issues/160
+        if x.dtype in dh.uint_dtypes:
+            assert dh.is_int_dtype(res.dtype)  # sanity check
+    else:
+        ph.assert_dtype("trace", in_dtype=x.dtype, out_dtype=res.dtype, expected=expected_dtype)
 
     n, m = x.shape[-2:]
     ph.assert_result_shape('trace', x.shape, res.shape, expected=x.shape[:-2])
diff --git a/array_api_tests/test_statistical_functions.py b/array_api_tests/test_statistical_functions.py
index 05d64e55..c8cdd573 100644
--- a/array_api_tests/test_statistical_functions.py
+++ b/array_api_tests/test_statistical_functions.py
@@ -130,44 +130,15 @@ def test_prod(x, data):
         out = xp.prod(x, **kw)
 
     dtype = kw.get("dtype", None)
-    if dtype is None:
-        if dh.is_int_dtype(x.dtype):
-            if x.dtype in dh.uint_dtypes:
-                default_dtype = dh.default_uint
-            else:
-                default_dtype = dh.default_int
-            if default_dtype is None:
-                _dtype = None
-            else:
-                m, M = dh.dtype_ranges[x.dtype]
-                d_m, d_M = dh.dtype_ranges[default_dtype]
-                if m < d_m or M > d_M:
-                    _dtype = x.dtype
-                else:
-                    _dtype = default_dtype
-        elif dh.is_float_dtype(x.dtype, include_complex=False):
-            if dh.dtype_nbits[x.dtype] > dh.dtype_nbits[dh.default_float]:
-                _dtype = x.dtype
-            else:
-                _dtype = dh.default_float
-        elif api_version > "2021.12":
-            # Complex dtype
-            if dh.dtype_nbits[x.dtype] > dh.dtype_nbits[dh.default_complex]:
-                _dtype = x.dtype
-            else:
-                _dtype = dh.default_complex
-        else:
-            raise RuntimeError("Unexpected dtype. This indicates a bug in the test suite.")
-    else:
-        _dtype = dtype
-    if _dtype is None:
+    expected_dtype = dh.accumulation_result_dtype(x.dtype, dtype)
+    if expected_dtype is None:
         # If a default uint cannot exist (i.e. in PyTorch which doesn't support
         # uint32 or uint64), we skip testing the output dtype.
         # See https://github.com/data-apis/array-api-tests/issues/106
         if x.dtype in dh.uint_dtypes:
             assert dh.is_int_dtype(out.dtype)  # sanity check
     else:
-        ph.assert_dtype("prod", in_dtype=x.dtype, out_dtype=out.dtype, expected=_dtype)
+        ph.assert_dtype("prod", in_dtype=x.dtype, out_dtype=out.dtype, expected=expected_dtype)
     _axes = sh.normalise_axis(kw.get("axis", None), x.ndim)
     ph.assert_keepdimable_shape(
         "prod", in_shape=x.shape, out_shape=out.shape, axes=_axes, keepdims=keepdims, kw=kw
@@ -246,44 +217,15 @@ def test_sum(x, data):
         out = xp.sum(x, **kw)
 
     dtype = kw.get("dtype", None)
-    if dtype is None:
-        if dh.is_int_dtype(x.dtype):
-            if x.dtype in dh.uint_dtypes:
-                default_dtype = dh.default_uint
-            else:
-                default_dtype = dh.default_int
-            if default_dtype is None:
-                _dtype = None
-            else:
-                m, M = dh.dtype_ranges[x.dtype]
-                d_m, d_M = dh.dtype_ranges[default_dtype]
-                if m < d_m or M > d_M:
-                    _dtype = x.dtype
-                else:
-                    _dtype = default_dtype
-        elif dh.is_float_dtype(x.dtype, include_complex=False):
-            if dh.dtype_nbits[x.dtype] > dh.dtype_nbits[dh.default_float]:
-                _dtype = x.dtype
-            else:
-                _dtype = dh.default_float
-        elif api_version > "2021.12":
-            # Complex dtype
-            if dh.dtype_nbits[x.dtype] > dh.dtype_nbits[dh.default_complex]:
-                _dtype = x.dtype
-            else:
-                _dtype = dh.default_complex
-        else:
-            raise RuntimeError("Unexpected dtype. This indicates a bug in the test suite.")
-    else:
-        _dtype = dtype
-    if _dtype is None:
+    expected_dtype = dh.accumulation_result_dtype(x.dtype, dtype)
+    if expected_dtype is None:
         # If a default uint cannot exist (i.e. in PyTorch which doesn't support
         # uint32 or uint64), we skip testing the output dtype.
         # See https://github.com/data-apis/array-api-tests/issues/160
         if x.dtype in dh.uint_dtypes:
             assert dh.is_int_dtype(out.dtype)  # sanity check
     else:
-        ph.assert_dtype("sum", in_dtype=x.dtype, out_dtype=out.dtype, expected=_dtype)
+        ph.assert_dtype("sum", in_dtype=x.dtype, out_dtype=out.dtype, expected=expected_dtype)
     _axes = sh.normalise_axis(kw.get("axis", None), x.ndim)
     ph.assert_keepdimable_shape(
         "sum", in_shape=x.shape, out_shape=out.shape, axes=_axes, keepdims=keepdims, kw=kw

From 5aa90835381618a736d300d2eae3351974a65806 Mon Sep 17 00:00:00 2001
From: Aaron Meurer <asmeurer@gmail.com>
Date: Wed, 7 Feb 2024 15:23:15 -0700
Subject: [PATCH 47/54] Remove unused allclose and assert_allclose helpers

---
 array_api_tests/array_helpers.py | 41 --------------------------------
 1 file changed, 41 deletions(-)

diff --git a/array_api_tests/array_helpers.py b/array_api_tests/array_helpers.py
index 9775ddac..c11d5732 100644
--- a/array_api_tests/array_helpers.py
+++ b/array_api_tests/array_helpers.py
@@ -10,10 +10,6 @@
 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',
            'where', 'isfinite', 'equal', 'not_equal', 'zero', 'one', 'NaN',
@@ -150,43 +146,6 @@ 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.

From 938f08698b72323e0943ce25f8d920ed1f14e9d0 Mon Sep 17 00:00:00 2001
From: Aaron Meurer <asmeurer@gmail.com>
Date: Thu, 15 Feb 2024 18:01:38 -0700
Subject: [PATCH 48/54] Update ndindex version requirement

---
 requirements.txt | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/requirements.txt b/requirements.txt
index bb33bc90..7e898020 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,4 +1,4 @@
 pytest
 pytest-json-report
 hypothesis>=6.68.0
-ndindex>=1.6
+ndindex>=1.8

From 3856b8fb6a8dace6afc83fa7bccad7629b4b61e7 Mon Sep 17 00:00:00 2001
From: Aaron Meurer <asmeurer@gmail.com>
Date: Thu, 15 Feb 2024 18:03:49 -0700
Subject: [PATCH 49/54] Fix linting issue

---
 array_api_tests/test_statistical_functions.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/array_api_tests/test_statistical_functions.py b/array_api_tests/test_statistical_functions.py
index c8cdd573..28130b84 100644
--- a/array_api_tests/test_statistical_functions.py
+++ b/array_api_tests/test_statistical_functions.py
@@ -11,7 +11,7 @@
 from . import hypothesis_helpers as hh
 from . import pytest_helpers as ph
 from . import shape_helpers as sh
-from . import xps, api_version
+from . import xps
 from ._array_module import _UndefinedStub
 from .typing import DataType
 

From ccc6ca3ef1cf736e32a55c01ff4496e1a4ce9a2a Mon Sep 17 00:00:00 2001
From: Matthew Barber <quitesimplymatt@gmail.com>
Date: Tue, 20 Feb 2024 16:24:28 +0000
Subject: [PATCH 50/54] Skip `test_cross` in CI

---
 numpy-skips.txt | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/numpy-skips.txt b/numpy-skips.txt
index 0c6f39ae..aebc249d 100644
--- a/numpy-skips.txt
+++ b/numpy-skips.txt
@@ -14,6 +14,8 @@ array_api_tests/test_constants.py::test_newaxis
 
 # linalg.solve issue in numpy.array_api as of v1.26.2 (see numpy#25146)
 array_api_tests/test_linalg.py::test_solve
+# numpy.array_api needs updating... or replaced on CI
+array_api_tests/test_linalg.py::test_cross
 
 # https://github.com/numpy/numpy/issues/21373
 array_api_tests/test_array_object.py::test_getitem

From 30924228111892cfa10f53544554a44747a5c245 Mon Sep 17 00:00:00 2001
From: Aaron Meurer <asmeurer@gmail.com>
Date: Thu, 22 Feb 2024 17:15:30 -0700
Subject: [PATCH 51/54] Test matmul, matrix_transpose, tensordot, and vecdot
 for the main and linalg namespaces separately

---
 array_api_tests/test_linalg.py | 105 ++++++++++++++++++++++++---------
 1 file changed, 76 insertions(+), 29 deletions(-)

diff --git a/array_api_tests/test_linalg.py b/array_api_tests/test_linalg.py
index 62862bba..35ff1d42 100644
--- a/array_api_tests/test_linalg.py
+++ b/array_api_tests/test_linalg.py
@@ -331,10 +331,9 @@ def test_inv(x):
 
     # TODO: Test that the result is actually the inverse
 
-@given(
-    *two_mutual_arrays(dh.real_dtypes)
-)
-def test_matmul(x1, x2):
+def _test_matmul(namespace, x1, x2):
+    matmul = namespace.matmul
+
     # TODO: Make this also test the @ operator
     if (x1.shape == () or x2.shape == ()
         or len(x1.shape) == len(x2.shape) == 1 and x1.shape != x2.shape
@@ -347,7 +346,7 @@ def test_matmul(x1, x2):
                "matmul did not raise an exception for invalid shapes")
         return
     else:
-        res = _array_module.matmul(x1, x2)
+        res = matmul(x1, x2)
 
     ph.assert_dtype("matmul", in_dtype=[x1.dtype, x2.dtype], out_dtype=res.dtype)
 
@@ -358,19 +357,32 @@ def test_matmul(x1, x2):
         ph.assert_result_shape("matmul", in_shapes=[x1.shape, x2.shape],
                                out_shape=res.shape,
                                expected=x2.shape[:-2] + x2.shape[-1:])
-        _test_stacks(_array_module.matmul, x1, x2, res=res, dims=1,
+        _test_stacks(matmul, x1, x2, res=res, dims=1,
                      matrix_axes=[(0,), (-2, -1)], res_axes=[-1])
     elif len(x2.shape) == 1:
         ph.assert_result_shape("matmul", in_shapes=[x1.shape, x2.shape],
                                out_shape=res.shape, expected=x1.shape[:-1])
-        _test_stacks(_array_module.matmul, x1, x2, res=res, dims=1,
+        _test_stacks(matmul, x1, x2, res=res, dims=1,
                      matrix_axes=[(-2, -1), (0,)], res_axes=[-1])
     else:
         stack_shape = sh.broadcast_shapes(x1.shape[:-2], x2.shape[:-2])
         ph.assert_result_shape("matmul", in_shapes=[x1.shape, x2.shape],
                                out_shape=res.shape,
                                expected=stack_shape + (x1.shape[-2], x2.shape[-1]))
-        _test_stacks(_array_module.matmul, x1, x2, res=res)
+        _test_stacks(matmul, x1, x2, res=res)
+
+@pytest.mark.xp_extension('linalg')
+@given(
+    *two_mutual_arrays(dh.real_dtypes)
+)
+def test_linalg_matmul(x1, x2):
+    return _test_matmul(linalg, x1, x2)
+
+@given(
+    *two_mutual_arrays(dh.real_dtypes)
+)
+def test_matmul(x1, x2):
+    return _test_matmul(_array_module, x1, x2)
 
 @pytest.mark.xp_extension('linalg')
 @given(
@@ -428,11 +440,9 @@ def test_matrix_power(x, n):
 def test_matrix_rank(x, kw):
     linalg.matrix_rank(x, **kw)
 
-@given(
-    x=arrays(dtype=xps.scalar_dtypes(), shape=matrix_shapes()),
-)
-def test_matrix_transpose(x):
-    res = _array_module.matrix_transpose(x)
+def _test_matrix_transpose(namespace, x):
+    matrix_transpose = namespace.matrix_transpose
+    res = matrix_transpose(x)
     true_val = lambda a: _array_module.asarray([[a[i, j] for i in
                                                 range(a.shape[0])] for j in
                                                 range(a.shape[1])],
@@ -444,7 +454,20 @@ def test_matrix_transpose(x):
     ph.assert_result_shape("matrix_transpose", in_shapes=[x.shape],
                            out_shape=res.shape, expected=shape)
 
-    _test_stacks(_array_module.matrix_transpose, x, res=res, true_val=true_val)
+    _test_stacks(matrix_transpose, x, res=res, true_val=true_val)
+
+@pytest.mark.xp_extension('linalg')
+@given(
+    x=arrays(dtype=xps.scalar_dtypes(), shape=matrix_shapes()),
+)
+def test_linalg_matrix_transpose(x):
+    return _test_matrix_transpose(linalg, x)
+
+@given(
+    x=arrays(dtype=xps.scalar_dtypes(), shape=matrix_shapes()),
+)
+def test_matrix_transpose(x):
+    return _test_matrix_transpose(_array_module, x)
 
 @pytest.mark.xp_extension('linalg')
 @given(
@@ -759,12 +782,9 @@ def _test_tensordot_stacks(x1, x2, kw, res):
         decomp_res_stack = xp.tensordot(x1_stack, x2_stack, axes=res_axes)
         assert_equal(res_stack, decomp_res_stack)
 
-@given(
-    *two_mutual_arrays(dh.numeric_dtypes, two_shapes=tensordot_shapes()),
-    tensordot_kw,
-)
-def test_tensordot(x1, x2, kw):
-    res = xp.tensordot(x1, x2, **kw)
+def _test_tensordot(namespace, x1, x2, kw):
+    tensordot = namespace.tensordot
+    res = tensordot(x1, x2, **kw)
 
     ph.assert_dtype("tensordot", in_dtype=[x1.dtype, x2.dtype],
                     out_dtype=res.dtype)
@@ -785,6 +805,21 @@ def test_tensordot(x1, x2, kw):
                            expected=result_shape)
     _test_tensordot_stacks(x1, x2, kw, res)
 
+@pytest.mark.xp_extension('linalg')
+@given(
+    *two_mutual_arrays(dh.numeric_dtypes, two_shapes=tensordot_shapes()),
+    tensordot_kw,
+)
+def test_linalg_tensordot(x1, x2, kw):
+    _test_tensordot(linalg, x1, x2, kw)
+
+@given(
+    *two_mutual_arrays(dh.numeric_dtypes, two_shapes=tensordot_shapes()),
+    tensordot_kw,
+)
+def test_tensordot(x1, x2, kw):
+    _test_tensordot(_array_module, x1, x2, kw)
+
 @pytest.mark.xp_extension('linalg')
 @given(
     x=arrays(dtype=xps.numeric_dtypes(), shape=matrix_shapes()),
@@ -828,12 +863,8 @@ def true_trace(x_stack, offset=0):
 
     _test_stacks(linalg.trace, x, **kw, res=res, dims=0, true_val=true_trace)
 
-
-@given(
-    *two_mutual_arrays(dh.numeric_dtypes, mutually_broadcastable_shapes(2, min_dims=1)),
-    data(),
-)
-def test_vecdot(x1, x2, data):
+def _test_vecdot(namespace, x1, x2, data):
+    vecdot = namespace.vecdot
     broadcasted_shape = sh.broadcast_shapes(x1.shape, x2.shape)
     min_ndim = min(x1.ndim, x2.ndim)
     ndim = len(broadcasted_shape)
@@ -842,14 +873,14 @@ def test_vecdot(x1, x2, data):
     x1_shape = (1,)*(ndim - x1.ndim) + tuple(x1.shape)
     x2_shape = (1,)*(ndim - x2.ndim) + tuple(x2.shape)
     if x1_shape[axis] != x2_shape[axis]:
-        ph.raises(Exception, lambda: xp.vecdot(x1, x2, **kw),
+        ph.raises(Exception, lambda: vecdot(x1, x2, **kw),
                   "vecdot did not raise an exception for invalid shapes")
         return
     expected_shape = list(broadcasted_shape)
     expected_shape.pop(axis)
     expected_shape = tuple(expected_shape)
 
-    res = xp.vecdot(x1, x2, **kw)
+    res = vecdot(x1, x2, **kw)
 
     ph.assert_dtype("vecdot", in_dtype=[x1.dtype, x2.dtype],
                     out_dtype=res.dtype)
@@ -862,9 +893,25 @@ def true_val(x, y, axis=-1):
     else:
         true_val = None
 
-    _test_stacks(linalg.vecdot, x1, x2, res=res, dims=0,
+    _test_stacks(vecdot, x1, x2, res=res, dims=0,
                  matrix_axes=(axis,), true_val=true_val)
 
+
+@pytest.mark.xp_extension('linalg')
+@given(
+    *two_mutual_arrays(dh.numeric_dtypes, mutually_broadcastable_shapes(2, min_dims=1)),
+    data(),
+)
+def test_linalg_vecdot(x1, x2, data):
+    _test_vecdot(linalg, x1, x2, data)
+
+@given(
+    *two_mutual_arrays(dh.numeric_dtypes, mutually_broadcastable_shapes(2, min_dims=1)),
+    data(),
+)
+def test_vecdot(x1, x2, data):
+    _test_vecdot(_array_module, x1, x2, data)
+
 # 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

From 3fefd20a78900442ac30cb12873177aab8616650 Mon Sep 17 00:00:00 2001
From: Matthew Barber <quitesimplymatt@gmail.com>
Date: Mon, 26 Feb 2024 10:53:11 +0000
Subject: [PATCH 52/54] Remove need for filtering in `invertible_matrices()`

---
 array_api_tests/hypothesis_helpers.py | 12 +++++++++---
 1 file changed, 9 insertions(+), 3 deletions(-)

diff --git a/array_api_tests/hypothesis_helpers.py b/array_api_tests/hypothesis_helpers.py
index b963bd7f..59ae5470 100644
--- a/array_api_tests/hypothesis_helpers.py
+++ b/array_api_tests/hypothesis_helpers.py
@@ -313,12 +313,18 @@ def invertible_matrices(draw, dtypes=xps.floating_dtypes(), stack_shapes=shapes(
     # For now, just generate stacks of diagonal matrices.
     n = draw(integers(0, SQRT_MAX_ARRAY_SIZE),)
     stack_shape = draw(stack_shapes)
-    d = draw(arrays(dtypes, shape=(*stack_shape, 1, n),
-                        elements=dict(allow_nan=False, allow_infinity=False)))
+    dtype = draw(dtypes)
+    elements = one_of(
+        from_dtype(dtype, min_value=0.5, allow_nan=False, allow_infinity=False),
+        from_dtype(dtype, max_value=-0.5, allow_nan=False, allow_infinity=False),
+    )
+    d = draw(arrays(dtype, shape=(*stack_shape, 1, n), elements=elements))
+
     # Functions that require invertible matrices may do anything when it is
     # singular, including raising an exception, so we make sure the diagonals
     # are sufficiently nonzero to avoid any numerical issues.
-    assume(xp.all(xp.abs(d) > 0.5))
+    assert xp.all(xp.abs(d) >= 0.5)
+
     diag_mask = xp.arange(n) == xp.reshape(xp.arange(n), (n, 1))
     return xp.where(diag_mask, d, xp.zeros_like(d))
 

From 268682dc233dda5e43efd0429e728efba66f9a77 Mon Sep 17 00:00:00 2001
From: Matthew Barber <quitesimplymatt@gmail.com>
Date: Mon, 26 Feb 2024 13:23:29 +0000
Subject: [PATCH 53/54] Skip flaky `test_reshape`

---
 array_api_tests/test_manipulation_functions.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/array_api_tests/test_manipulation_functions.py b/array_api_tests/test_manipulation_functions.py
index fd27b2dc..8cbe7750 100644
--- a/array_api_tests/test_manipulation_functions.py
+++ b/array_api_tests/test_manipulation_functions.py
@@ -250,6 +250,7 @@ def reshape_shapes(draw, shape):
     return tuple(rshape)
 
 
+@pytest.mark.skip("flaky")  # TODO: fix!
 @given(
     x=hh.arrays(dtype=xps.scalar_dtypes(), shape=hh.shapes(max_side=MAX_SIDE)),
     data=st.data(),

From 0ddb0cd511e8f9c72bc29dca0e27e188cb24021e Mon Sep 17 00:00:00 2001
From: Matthew Barber <quitesimplymatt@gmail.com>
Date: Mon, 26 Feb 2024 14:33:08 +0000
Subject: [PATCH 54/54] Less filtering in `positive_definitive_matrices`

---
 array_api_tests/hypothesis_helpers.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/array_api_tests/hypothesis_helpers.py b/array_api_tests/hypothesis_helpers.py
index 59ae5470..3033dac3 100644
--- a/array_api_tests/hypothesis_helpers.py
+++ b/array_api_tests/hypothesis_helpers.py
@@ -302,8 +302,9 @@ def positive_definite_matrices(draw, dtypes=xps.floating_dtypes()):
     # TODO: Generate arbitrary positive definite matrices, for instance, by
     # using something like
     # https://github.com/scikit-learn/scikit-learn/blob/844b4be24/sklearn/datasets/_samples_generator.py#L1351.
-    n = draw(integers(0))
-    shape = draw(shapes()) + (n, n)
+    base_shape = draw(shapes())
+    n = draw(integers(0, 8))  # 8 is an arbitrary small but interesting-enough value
+    shape = base_shape + (n, n)
     assume(prod(i for i in shape if i) < MAX_ARRAY_SIZE)
     dtype = draw(dtypes)
     return broadcast_to(eye(n, dtype=dtype), shape)