diff --git a/array_api_tests/test_linalg.py b/array_api_tests/test_linalg.py index 29d7775b..acc4dc42 100644 --- a/array_api_tests/test_linalg.py +++ b/array_api_tests/test_linalg.py @@ -516,10 +516,12 @@ def test_matrix_rank(x, kw): 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])], - dtype=a.dtype) + def true_val(a): + if 0 in a.shape: + return xp.empty((a.shape[1], a.shape[0]), + dtype=a.dtype, device=a.device) + return xp.stack([a[i, :] for i in range(a.shape[0])], axis=1) + shape = list(x.shape) shape[-1], shape[-2] = shape[-2], shape[-1] shape = tuple(shape) diff --git a/meta_tests/test_linalg.py b/meta_tests/test_linalg.py index 82794b6c..1cdd2860 100644 --- a/meta_tests/test_linalg.py +++ b/meta_tests/test_linalg.py @@ -5,6 +5,31 @@ from array_api_tests .hypothesis_helpers import symmetric_matrices from array_api_tests import array_helpers as ah from array_api_tests import _array_module as xp +from array_api_tests.test_linalg import _test_matrix_transpose + + +@pytest.mark.parametrize('shape', [(2, 3), (2, 2, 3), (0, 3), (3, 0), (0, 0)]) +def test_matrix_transpose_without_nested_arrays(monkeypatch, shape): + size = 1 + for dim in shape: + size *= dim + x = xp.reshape(xp.arange(size, dtype=xp.int64), shape) + original_asarray = xp.asarray + + def asarray(obj, **kwargs): + def check_sequence(value): + if isinstance(value, (list, tuple)): + for item in value: + check_sequence(item) + else: + assert isinstance(value, (bool, int, float, complex)) + + if isinstance(obj, (list, tuple)): + check_sequence(obj) + return original_asarray(obj, **kwargs) + + monkeypatch.setattr(xp, 'asarray', asarray) + _test_matrix_transpose(xp, x) @pytest.mark.xp_extension('linalg') @given(x=symmetric_matrices(finite=True))