scipy 1.15.3__cp312-cp312-win_amd64.whl → 1.16.0rc2__cp312-cp312-win_amd64.whl

scipy/linalg/tests/test_batch.py

@@ -0,0 +1,588 @@
+import inspect
+import pytest
+import numpy as np
+from numpy.testing import assert_allclose
+from scipy import linalg, sparse
+
+
+real_floating = [np.float32, np.float64]
+complex_floating = [np.complex64, np.complex128]
+floating = real_floating + complex_floating
+
+
+def get_random(shape, *, dtype, rng):
+    A = rng.random(shape)
+    if np.issubdtype(dtype, np.complexfloating):
+        A = A + rng.random(shape) * 1j
+    return A.astype(dtype)
+
+def get_nearly_hermitian(shape, dtype, atol, rng):
+    # Generate a batch of nearly Hermitian matrices with specified
+    # `shape` and `dtype`. `atol` controls the level of noise in
+    # Hermitian-ness to by generated by `rng`.
+    A = rng.random(shape).astype(dtype)
+    At = np.conj(A.swapaxes(-1, -2))
+    noise = rng.standard_normal(size=A.shape).astype(dtype) * atol
+    return A + At + noise
+
+
+class TestBatch:
+    # Test batch support for most linalg functions
+
+    def batch_test(self, fun, arrays, *, core_dim=2, n_out=1, kwargs=None, dtype=None,
+                   broadcast=True, check_kwargs=True):
+        # Check that all outputs of batched call `fun(A, **kwargs)` are the same
+        # as if we loop over the separate vectors/matrices in `A`. Also check
+        # that `fun` accepts `A` by position or keyword and that results are
+        # identical. This is important because the name of the array argument
+        # is manually specified to the decorator, and it's easy to mess up.
+        # However, this makes it hard to test positional arguments passed
+        # after the array, so we test that separately for a few functions to
+        # make sure the decorator is working as it should.
+
+        kwargs = {} if kwargs is None else kwargs
+        parameters = list(inspect.signature(fun).parameters.keys())
+        arrays = (arrays,) if not isinstance(arrays, tuple) else arrays
+
+        # Identical results when passing argument by keyword or position
+        res2 = fun(*arrays, **kwargs)
+        if check_kwargs:
+            res1 = fun(**dict(zip(parameters, arrays)), **kwargs)
+            for out1, out2 in zip(res1, res2):  # even a single array is iterable...
+                np.testing.assert_equal(out1, out2)
+
+        # Check results vs looping over
+        res = (res2,) if n_out == 1 else res2
+        # This is not the general behavior (only batch dimensions get
+        # broadcasted by the decorator) but it's easier for testing.
+        if broadcast:
+            arrays = np.broadcast_arrays(*arrays)
+        batch_shape = arrays[0].shape[:-core_dim]
+        for i in range(batch_shape[0]):
+            for j in range(batch_shape[1]):
+                arrays_ij = (array[i, j] for array in arrays)
+                ref = fun(*arrays_ij, **kwargs)
+                ref = ((np.asarray(ref),) if n_out == 1 else
+                       tuple(np.asarray(refk) for refk in ref))
+                for k in range(n_out):
+                    assert_allclose(res[k][i, j], ref[k])
+                    assert np.shape(res[k][i, j]) == ref[k].shape
+
+        for k in range(len(ref)):
+            out_dtype = ref[k].dtype if dtype is None else dtype
+            assert res[k].dtype == out_dtype
+
+        return res2  # return original, non-tuplized result
+
+    @pytest.mark.parametrize('dtype', floating)
+    def test_expm_cond(self, dtype):
+        rng = np.random.default_rng(8342310302941288912051)
+        A = rng.random((5, 3, 4, 4)).astype(dtype)
+        self.batch_test(linalg.expm_cond, A)
+
+    @pytest.mark.parametrize('dtype', floating)
+    def test_issymmetric(self, dtype):
+        rng = np.random.default_rng(8342310302941288912051)
+        A = get_nearly_hermitian((5, 3, 4, 4), dtype, 3e-4, rng)
+        res = self.batch_test(linalg.issymmetric, A, kwargs=dict(atol=1e-3))
+        assert not np.all(res)  # ensure test is not trivial: not all True or False;
+        assert np.any(res)      # also confirms that `atol` is passed to issymmetric
+
+    @pytest.mark.parametrize('dtype', floating)
+    def test_ishermitian(self, dtype):
+        rng = np.random.default_rng(8342310302941288912051)
+        A = get_nearly_hermitian((5, 3, 4, 4), dtype, 3e-4, rng)
+        res = self.batch_test(linalg.ishermitian, A, kwargs=dict(atol=1e-3))
+        assert not np.all(res)  # ensure test is not trivial: not all True or False;
+        assert np.any(res)      # also confirms that `atol` is passed to ishermitian
+
+    @pytest.mark.parametrize('dtype', floating)
+    def test_diagsvd(self, dtype):
+        rng = np.random.default_rng(8342310302941288912051)
+        A = rng.random((5, 3, 4)).astype(dtype)
+        res1 = self.batch_test(linalg.diagsvd, A, kwargs=dict(M=6, N=4), core_dim=1)
+        # test that `M, N` can be passed by position
+        res2 = linalg.diagsvd(A, 6, 4)
+        np.testing.assert_equal(res1, res2)
+
+    @pytest.mark.parametrize('fun', [linalg.inv, linalg.sqrtm, linalg.signm,
+                                     linalg.sinm, linalg.cosm, linalg.tanhm,
+                                     linalg.sinhm, linalg.coshm, linalg.tanhm,
+                                     linalg.pinv, linalg.pinvh, linalg.orth])
+    @pytest.mark.parametrize('dtype', floating)
+    def test_matmat(self, fun, dtype):  # matrix in, matrix out
+        rng = np.random.default_rng(8342310302941288912051)
+        A = get_random((5, 3, 4, 4), dtype=dtype, rng=rng)
+
+        # sqrtm can return complex output for real input resulting in i/o type
+        # mismatch. Nudge the eigenvalues to positive side to avoid this.
+        if fun == linalg.sqrtm:
+            A = A + 3*np.eye(4, dtype=dtype)
+
+        self.batch_test(fun, A)
+
+    @pytest.mark.parametrize('dtype', floating)
+    def test_null_space(self, dtype):
+        rng = np.random.default_rng(8342310302941288912051)
+        A = get_random((5, 3, 4, 6), dtype=dtype, rng=rng)
+        self.batch_test(linalg.null_space, A)
+
+    @pytest.mark.parametrize('dtype', floating)
+    def test_funm(self, dtype):
+        rng = np.random.default_rng(8342310302941288912051)
+        A = get_random((2, 4, 3, 3), dtype=dtype, rng=rng)
+        self.batch_test(linalg.funm, A, kwargs=dict(func=np.sin))
+
+    @pytest.mark.parametrize('dtype', floating)
+    def test_fractional_matrix_power(self, dtype):
+        rng = np.random.default_rng(8342310302941288912051)
+        A = get_random((2, 4, 3, 3), dtype=dtype, rng=rng)
+        res1 = self.batch_test(linalg.fractional_matrix_power, A, kwargs={'t':1.5})
+        # test that `t` can be passed by position
+        res2 = linalg.fractional_matrix_power(A, 1.5)
+        np.testing.assert_equal(res1, res2)
+
+    @pytest.mark.parametrize('dtype', floating)
+    def test_logm(self, dtype):
+        # One test failed absolute tolerance with default random seed
+        rng = np.random.default_rng(89940026998903887141749720079406074936)
+        A = get_random((5, 3, 4, 4), dtype=dtype, rng=rng)
+        A = A + 3*np.eye(4)  # avoid complex output for real input
+        res1 = self.batch_test(linalg.logm, A)
+        # test that `disp` can be passed by position
+        res2 = linalg.logm(A)
+        for res1i, res2i in zip(res1, res2):
+            np.testing.assert_equal(res1i, res2i)
+
+    @pytest.mark.parametrize('dtype', floating)
+    def test_pinv(self, dtype):
+        rng = np.random.default_rng(8342310302941288912051)
+        A = get_random((5, 3, 4, 4), dtype=dtype, rng=rng)
+        self.batch_test(linalg.pinv, A, n_out=2, kwargs=dict(return_rank=True))
+
+    @pytest.mark.parametrize('dtype', floating)
+    def test_matrix_balance(self, dtype):
+        rng = np.random.default_rng(8342310302941288912051)
+        A = get_random((5, 3, 4, 4), dtype=dtype, rng=rng)
+        self.batch_test(linalg.matrix_balance, A, n_out=2)
+        self.batch_test(linalg.matrix_balance, A, n_out=2, kwargs={'separate':True})
+
+    @pytest.mark.parametrize('dtype', floating)
+    def test_bandwidth(self, dtype):
+        rng = np.random.default_rng(8342310302941288912051)
+        A = get_random((4, 4), dtype=dtype, rng=rng)
+        A = np.asarray([np.triu(A, k) for k in range(-3, 3)]).reshape((2, 3, 4, 4))
+        self.batch_test(linalg.bandwidth, A, n_out=2)
+
+    @pytest.mark.parametrize('fun_n_out', [(linalg.cholesky, 1), (linalg.ldl, 3),
+                                           (linalg.cho_factor, 2)])
+    @pytest.mark.parametrize('dtype', floating)
+    def test_ldl_cholesky(self, fun_n_out, dtype):
+        rng = np.random.default_rng(8342310302941288912051)
+        fun, n_out = fun_n_out
+        A = get_nearly_hermitian((5, 3, 4, 4), dtype, 0, rng)  # exactly Hermitian
+        A = A + 4*np.eye(4, dtype=dtype)  # ensure positive definite for Cholesky
+        self.batch_test(fun, A, n_out=n_out)
+
+    @pytest.mark.parametrize('compute_uv', [False, True])
+    @pytest.mark.parametrize('dtype', floating)
+    def test_svd(self, compute_uv, dtype):
+        rng = np.random.default_rng(8342310302941288912051)
+        A = get_random((5, 3, 2, 4), dtype=dtype, rng=rng)
+        n_out = 3 if compute_uv else 1
+        self.batch_test(linalg.svd, A, n_out=n_out, kwargs=dict(compute_uv=compute_uv))
+
+    @pytest.mark.parametrize('fun', [linalg.polar, linalg.qr, linalg.rq])
+    @pytest.mark.parametrize('dtype', floating)
+    def test_polar_qr_rq(self, fun, dtype):
+        rng = np.random.default_rng(8342310302941288912051)
+        A = get_random((5, 3, 2, 4), dtype=dtype, rng=rng)
+        self.batch_test(fun, A, n_out=2)
+
+    @pytest.mark.parametrize('cdim', [(5,), (5, 4), (2, 3, 5, 4)])
+    @pytest.mark.parametrize('dtype', floating)
+    def test_qr_multiply(self, cdim, dtype):
+        rng = np.random.default_rng(8342310302941288912051)
+        A = get_random((2, 3, 5, 5), dtype=dtype, rng=rng)
+        c = get_random(cdim, dtype=dtype, rng=rng)
+        res = linalg.qr_multiply(A, c, mode='left')
+        q, r = linalg.qr(A)
+        ref = q @ c
+        atol = 1e-6 if dtype in {np.float32, np.complex64} else 1e-12
+        assert_allclose(res[0], ref, atol=atol)
+        assert_allclose(res[1], r, atol=atol)
+
+    @pytest.mark.parametrize('uvdim', [[(5,), (3,)], [(4, 5, 2), (4, 3, 2)]])
+    @pytest.mark.parametrize('dtype', floating)
+    def test_qr_update(self, uvdim, dtype):
+        rng = np.random.default_rng(8342310302941288912051)
+        udim, vdim = uvdim
+        A = get_random((4, 5, 3), dtype=dtype, rng=rng)
+        u = get_random(udim, dtype=dtype, rng=rng)
+        v = get_random(vdim, dtype=dtype, rng=rng)
+        q, r = linalg.qr(A)
+        res = linalg.qr_update(q, r, u, v)
+        for i in range(4):
+            qi, ri = q[i], r[i]
+            ui, vi = (u, v) if u.ndim == 1 else (u[i], v[i])
+            ref_i = linalg.qr_update(qi, ri, ui, vi)
+            assert_allclose(res[0][i], ref_i[0])
+            assert_allclose(res[1][i], ref_i[1])
+
+    @pytest.mark.parametrize('udim', [(5,), (4, 3, 5)])
+    @pytest.mark.parametrize('kdim', [(), (4,)])
+    @pytest.mark.parametrize('dtype', floating)
+    def test_qr_insert(self, udim, kdim, dtype):
+        rng = np.random.default_rng(8342310302941288912051)
+        A = get_random((4, 5, 5), dtype=dtype, rng=rng)
+        u = get_random(udim, dtype=dtype, rng=rng)
+        k = rng.integers(0, 5, size=kdim)
+        q, r = linalg.qr(A)
+        res = linalg.qr_insert(q, r, u, k)
+        for i in range(4):
+            qi, ri = q[i], r[i]
+            ki = k if k.ndim == 0 else k[i]
+            ui = u if u.ndim == 1 else u[i]
+            ref_i = linalg.qr_insert(qi, ri, ui, ki)
+            assert_allclose(res[0][i], ref_i[0])
+            assert_allclose(res[1][i], ref_i[1])
+
+    @pytest.mark.parametrize('kdim', [(), (4,)])
+    @pytest.mark.parametrize('dtype', floating)
+    def test_qr_delete(self, kdim, dtype):
+        rng = np.random.default_rng(8342310302941288912051)
+        A = get_random((4, 5, 5), dtype=dtype, rng=rng)
+        k = rng.integers(0, 4, size=kdim)
+        q, r = linalg.qr(A)
+        res = linalg.qr_delete(q, r, k)
+        for i in range(4):
+            qi, ri = q[i], r[i]
+            ki = k if k.ndim == 0 else k[i]
+            ref_i = linalg.qr_delete(qi, ri, ki)
+            assert_allclose(res[0][i], ref_i[0])
+            assert_allclose(res[1][i], ref_i[1])
+
+    @pytest.mark.parametrize('fun', [linalg.schur, linalg.lu_factor])
+    @pytest.mark.parametrize('dtype', floating)
+    def test_schur_lu(self, fun, dtype):
+        rng = np.random.default_rng(8342310302941288912051)
+        A = get_random((5, 3, 4, 4), dtype=dtype, rng=rng)
+        self.batch_test(fun, A, n_out=2)
+
+    @pytest.mark.parametrize('calc_q', [False, True])
+    @pytest.mark.parametrize('dtype', floating)
+    def test_hessenberg(self, calc_q, dtype):
+        rng = np.random.default_rng(8342310302941288912051)
+        A = get_random((5, 3, 4, 4), dtype=dtype, rng=rng)
+        n_out = 2 if calc_q else 1
+        self.batch_test(linalg.hessenberg, A, n_out=n_out, kwargs=dict(calc_q=calc_q))
+
+    @pytest.mark.parametrize('eigvals_only', [False, True])
+    @pytest.mark.parametrize('dtype', floating)
+    def test_eig_banded(self, eigvals_only, dtype):
+        rng = np.random.default_rng(8342310302941288912051)
+        A = get_random((5, 3, 4, 4), dtype=dtype, rng=rng)
+        n_out = 1 if eigvals_only else 2
+        self.batch_test(linalg.eig_banded, A, n_out=n_out,
+                        kwargs=dict(eigvals_only=eigvals_only))
+
+    @pytest.mark.parametrize('dtype', floating)
+    def test_eigvals_banded(self, dtype):
+        rng = np.random.default_rng(8342310302941288912051)
+        A = get_random((5, 3, 4, 4), dtype=dtype, rng=rng)
+        self.batch_test(linalg.eigvals_banded, A)
+
+    @pytest.mark.parametrize('two_in', [False, True])
+    @pytest.mark.parametrize('fun_n_nout', [(linalg.eigh, 1), (linalg.eigh, 2),
+                                            (linalg.eigvalsh, 1), (linalg.eigvals, 1)])
+    @pytest.mark.parametrize('dtype', floating)
+    def test_eigh(self, two_in, fun_n_nout, dtype):
+        rng = np.random.default_rng(8342310302941288912051)
+        fun, n_out = fun_n_nout
+        A = get_nearly_hermitian((1, 3, 4, 4), dtype, 0, rng)  # exactly Hermitian
+        B = get_nearly_hermitian((2, 1, 4, 4), dtype, 0, rng)  # exactly Hermitian
+        B = B + 4*np.eye(4).astype(dtype)  # needs to be positive definite
+        args = (A, B) if two_in else (A,)
+        kwargs = dict(eigvals_only=True) if (n_out == 1 and fun==linalg.eigh) else {}
+        self.batch_test(fun, args, n_out=n_out, kwargs=kwargs)
+
+    @pytest.mark.parametrize('compute_expm', [False, True])
+    @pytest.mark.parametrize('dtype', floating)
+    def test_expm_frechet(self, compute_expm, dtype):
+        rng = np.random.default_rng(8342310302941288912051)
+        A = get_random((1, 3, 4, 4), dtype=dtype, rng=rng)
+        E = get_random((2, 1, 4, 4), dtype=dtype, rng=rng)
+        n_out = 2 if compute_expm else 1
+        self.batch_test(linalg.expm_frechet, (A, E), n_out=n_out,
+                        kwargs=dict(compute_expm=compute_expm))
+
+    @pytest.mark.parametrize('dtype', floating)
+    def test_subspace_angles(self, dtype):
+        rng = np.random.default_rng(8342310302941288912051)
+        A = get_random((1, 3, 4, 3), dtype=dtype, rng=rng)
+        B = get_random((2, 1, 4, 3), dtype=dtype, rng=rng)
+        self.batch_test(linalg.subspace_angles, (A, B))
+        # just to show that A and B don't need to be broadcastable
+        M, N, K = 4, 5, 3
+        A = get_random((1, 3, M, N), dtype=dtype, rng=rng)
+        B = get_random((2, 1, M, K), dtype=dtype, rng=rng)
+        assert linalg.subspace_angles(A, B).shape == (2, 3, min(N, K))
+
+    @pytest.mark.parametrize('fun', [linalg.svdvals])
+    @pytest.mark.parametrize('dtype', floating)
+    def test_svdvals(self, fun, dtype):
+        rng = np.random.default_rng(8342310302941288912051)
+        A = get_random((2, 3, 4, 5), dtype=dtype, rng=rng)
+        self.batch_test(fun, A)
+
+    @pytest.mark.parametrize('fun_n_out', [(linalg.orthogonal_procrustes, 2),
+                                           (linalg.khatri_rao, 1),
+                                           (linalg.solve_continuous_lyapunov, 1),
+                                           (linalg.solve_discrete_lyapunov, 1),
+                                           (linalg.qz, 4),
+                                           (linalg.ordqz, 6)])
+    @pytest.mark.parametrize('dtype', floating)
+    def test_two_generic_matrix_inputs(self, fun_n_out, dtype):
+        rng = np.random.default_rng(8342310302941288912051)
+        fun, n_out = fun_n_out
+        A = get_random((2, 3, 4, 4), dtype=dtype, rng=rng)
+        B = get_random((2, 3, 4, 4), dtype=dtype, rng=rng)
+        self.batch_test(fun, (A, B), n_out=n_out)
+
+    @pytest.mark.parametrize('dtype', floating)
+    def test_cossin(self, dtype):
+        rng = np.random.default_rng(8342310302941288912051)
+        p, q = 3, 4
+        X = get_random((2, 3, 10, 10), dtype=dtype, rng=rng)
+        x11, x12, x21, x22 = (X[..., :p, :q], X[..., :p, q:],
+                              X[..., p:, :q], X[..., p:, q:])
+        res = linalg.cossin(X, p, q)
+        ref = linalg.cossin((x11, x12, x21, x22))
+        for res_i, ref_i in zip(res, ref):
+            np.testing.assert_equal(res_i, ref_i)
+
+        for j in range(2):
+            for k in range(3):
+                ref_jk = linalg.cossin(X[j, k], p, q)
+                for res_i, ref_ijk in zip(res, ref_jk):
+                    np.testing.assert_equal(res_i[j, k], ref_ijk)
+
+    @pytest.mark.parametrize('dtype', floating)
+    def test_sylvester(self, dtype):
+        rng = np.random.default_rng(8342310302941288912051)
+        A = get_random((2, 3, 5, 5), dtype=dtype, rng=rng)
+        B = get_random((2, 3, 5, 5), dtype=dtype, rng=rng)
+        C = get_random((2, 3, 5, 5), dtype=dtype, rng=rng)
+        self.batch_test(linalg.solve_sylvester, (A, B, C))
+
+    @pytest.mark.parametrize('fun', [linalg.solve_continuous_are,
+                                     linalg.solve_discrete_are])
+    @pytest.mark.parametrize('dtype', floating)
+    def test_are(self, fun, dtype):
+        rng = np.random.default_rng(8342310302941288912051)
+        a = get_random((2, 3, 5, 5), dtype=dtype, rng=rng)
+        b = get_random((2, 3, 5, 5), dtype=dtype, rng=rng)
+        q = get_nearly_hermitian((2, 3, 5, 5), dtype=dtype, atol=0, rng=rng)
+        r = get_nearly_hermitian((2, 3, 5, 5), dtype=dtype, atol=0, rng=rng)
+        a = a + 5*np.eye(5)  # making these positive definite seems to help
+        b = b + 5*np.eye(5)
+        q = q + 5*np.eye(5)
+        r = r + 5*np.eye(5)
+        # can't easily generate valid random e, s
+        self.batch_test(fun, (a, b, q, r))
+
+    @pytest.mark.parametrize('dtype', floating)
+    def test_rsf2cs(self, dtype):
+        rng = np.random.default_rng(8342310302941288912051)
+        A = get_random((2, 3, 4, 4), dtype=dtype, rng=rng)
+        T, Z = linalg.schur(A)
+        self.batch_test(linalg.rsf2csf, (T, Z), n_out=2)
+
+    @pytest.mark.parametrize('dtype', floating)
+    def test_cholesky_banded(self, dtype):
+        rng = np.random.default_rng(8342310302941288912051)
+        ab = get_random((5, 4, 3, 6), dtype=dtype, rng=rng)
+        ab[..., -1, :] = 10  # make diagonal dominant
+        self.batch_test(linalg.cholesky_banded, ab)
+
+    @pytest.mark.parametrize('dtype', floating)
+    def test_block_diag(self, dtype):
+        rng = np.random.default_rng(8342310302941288912051)
+        a = get_random((1, 3, 1, 3), dtype=dtype, rng=rng)
+        b = get_random((2, 1, 3, 6), dtype=dtype, rng=rng)
+        c = get_random((1, 1, 3, 2), dtype=dtype, rng=rng)
+
+        # batch_test doesn't have the logic to broadcast just the batch shapes,
+        # so do it manually.
+        a2 = np.broadcast_to(a, (2, 3, 1, 3))
+        b2 = np.broadcast_to(b, (2, 3, 3, 6))
+        c2 = np.broadcast_to(c, (2, 3, 3, 2))
+        ref = self.batch_test(linalg.block_diag, (a2, b2, c2),
+                              check_kwargs=False, broadcast=False)
+
+        # Check that `block_diag` broadcasts the batch shapes as expected.
+        res = linalg.block_diag(a, b, c)
+        assert_allclose(res, ref)
+
+    @pytest.mark.parametrize('fun_n_out', [(linalg.eigh_tridiagonal, 2),
+                                           (linalg.eigvalsh_tridiagonal, 1)])
+    @pytest.mark.parametrize('dtype', real_floating)
+    # "Only real arrays currently supported"
+    def test_eigh_tridiagonal(self, fun_n_out, dtype):
+        rng = np.random.default_rng(8342310302941288912051)
+        fun, n_out = fun_n_out
+        d = get_random((3, 4, 5), dtype=dtype, rng=rng)
+        e = get_random((3, 4, 4), dtype=dtype, rng=rng)
+        self.batch_test(fun, (d, e), core_dim=1, n_out=n_out, broadcast=False)
+
+    @pytest.mark.parametrize('bdim', [(5,), (5, 4), (2, 3, 5, 4)])
+    @pytest.mark.parametrize('dtype', floating)
+    def test_solve(self, bdim, dtype):
+        rng = np.random.default_rng(8342310302941288912051)
+        A = get_random((2, 3, 5, 5), dtype=dtype, rng=rng)
+        b = get_random(bdim, dtype=dtype, rng=rng)
+        x = linalg.solve(A, b)
+        if len(bdim) == 1:
+            x = x[..., np.newaxis]
+            b = b[..., np.newaxis]
+        assert_allclose(A @ x - b, 0, atol=1.5e-6)
+        assert_allclose(x, np.linalg.solve(A, b), atol=3e-6)
+
+    @pytest.mark.parametrize('bdim', [(5,), (5, 4), (2, 3, 5, 4)])
+    @pytest.mark.parametrize('dtype', floating)
+    def test_lu_solve(self, bdim, dtype):
+        rng = np.random.default_rng(8342310302941288912051)
+        A = get_random((2, 3, 5, 5), dtype=dtype, rng=rng)
+        b = get_random(bdim, dtype=dtype, rng=rng)
+        lu_and_piv = linalg.lu_factor(A)
+        x = linalg.lu_solve(lu_and_piv, b)
+        if len(bdim) == 1:
+            x = x[..., np.newaxis]
+            b = b[..., np.newaxis]
+        assert_allclose(A @ x - b, 0, atol=1.5e-6)
+        assert_allclose(x, np.linalg.solve(A, b), atol=3e-6)
+
+    @pytest.mark.parametrize('l_and_u', [(1, 1), ([2, 1, 0], [0, 1 , 2])])
+    @pytest.mark.parametrize('bdim', [(5,), (5, 4), (2, 3, 5, 4)])
+    @pytest.mark.parametrize('dtype', floating)
+    def test_solve_banded(self, l_and_u, bdim, dtype):
+        rng = np.random.default_rng(8342310302941288912051)
+        l, u = l_and_u
+        ab = get_random((2, 3, 3, 5), dtype=dtype, rng=rng)
+        b = get_random(bdim, dtype=dtype, rng=rng)
+        x = linalg.solve_banded((l, u), ab, b)
+        for i in range(2):
+            for j in range(3):
+                bij = b if len(bdim) <= 2 else b[i, j]
+                lj = l if np.ndim(l) == 0 else l[j]
+                uj = u if np.ndim(u) == 0 else u[j]
+                xij = linalg.solve_banded((lj, uj), ab[i, j], bij)
+                assert_allclose(x[i, j], xij)
+
+    # Can uncomment when `solve_toeplitz` deprecation is done (SciPy 1.17)
+    # @pytest.mark.parametrize('separate_r', [False, True])
+    # @pytest.mark.parametrize('bdim', [(5,), (5, 4), (2, 3, 5, 4)])
+    # @pytest.mark.parametrize('dtype', floating)
+    # def test_solve_toeplitz(self, separate_r, bdim, dtype):
+    #     rng = np.random.default_rng(8342310302941288912051)
+    #     c = get_random((2, 3, 5), dtype=dtype, rng=rng)
+    #     r = get_random((2, 3, 5), dtype=dtype, rng=rng)
+    #     c_or_cr = (c, r) if separate_r else c
+    #     b = get_random(bdim, dtype=dtype, rng=rng)
+    #     x = linalg.solve_toeplitz(c_or_cr, b)
+    #     for i in range(2):
+    #         for j in range(3):
+    #             bij = b if len(bdim) <= 2 else b[i, j]
+    #             c_or_cr_ij = (c[i, j], r[i, j]) if separate_r else c[i, j]
+    #             xij = linalg.solve_toeplitz(c_or_cr_ij, bij)
+    #             assert_allclose(x[i, j], xij)
+
+    @pytest.mark.parametrize('bdim', [(5,), (5, 4), (2, 3, 5, 4)])
+    @pytest.mark.parametrize('dtype', floating)
+    def test_cho_solve(self, bdim, dtype):
+        rng = np.random.default_rng(8342310302941288912051)
+        A = get_nearly_hermitian((2, 3, 5, 5), dtype=dtype, atol=0, rng=rng)
+        A = A + 5*np.eye(5)
+        c_and_lower = linalg.cho_factor(A)
+        b = get_random(bdim, dtype=dtype, rng=rng)
+        x = linalg.cho_solve(c_and_lower, b)
+        if len(bdim) == 1:
+            x = x[..., np.newaxis]
+            b = b[..., np.newaxis]
+        assert_allclose(A @ x - b, 0, atol=1e-6)
+        assert_allclose(x, np.linalg.solve(A, b), atol=2e-6)
+
+    @pytest.mark.parametrize('lower', [False, True])
+    @pytest.mark.parametrize('bdim', [(5,), (5, 4), (2, 3, 5, 4)])
+    @pytest.mark.parametrize('dtype', floating)
+    def test_cho_solve_banded(self, lower, bdim, dtype):
+        rng = np.random.default_rng(8342310302941288912051)
+        A = get_random((2, 3, 3, 5), dtype=dtype, rng=rng)
+        row_diag = 0 if lower else -1
+        A[:, :, row_diag] = 10
+        cb = linalg.cholesky_banded(A, lower=lower)
+        b = get_random(bdim, dtype=dtype, rng=rng)
+        x = linalg.cho_solve_banded((cb, lower), b)
+        for i in range(2):
+            for j in range(3):
+                bij = b if len(bdim) <= 2 else b[i, j]
+                xij = linalg.cho_solve_banded((cb[i, j], lower), bij)
+                assert_allclose(x[i, j], xij)
+
+    @pytest.mark.parametrize('bdim', [(5,), (5, 4), (2, 3, 5, 4)])
+    @pytest.mark.parametrize('dtype', floating)
+    def test_solveh_banded(self, bdim, dtype):
+        rng = np.random.default_rng(8342310302941288912051)
+        A = get_random((2, 3, 3, 5), dtype=dtype, rng=rng)
+        A[:, :, -1] = 10
+        b = get_random(bdim, dtype=dtype, rng=rng)
+        x = linalg.solveh_banded(A, b)
+        for i in range(2):
+            for j in range(3):
+                bij = b if len(bdim) <= 2 else b[i, j]
+                xij = linalg.solveh_banded(A[i, j], bij)
+                assert_allclose(x[i, j], xij)
+
+    @pytest.mark.parametrize('bdim', [(5,), (5, 4), (2, 3, 5, 4)])
+    @pytest.mark.parametrize('dtype', floating)
+    def test_solve_triangular(self, bdim, dtype):
+        rng = np.random.default_rng(8342310302941288912051)
+        A = get_random((2, 3, 5, 5), dtype=dtype, rng=rng)
+        A = np.tril(A)
+        b = get_random(bdim, dtype=dtype, rng=rng)
+        x = linalg.solve_triangular(A, b, lower=True)
+        if len(bdim) == 1:
+            x = x[..., np.newaxis]
+            b = b[..., np.newaxis]
+        atol = 1e-10 if dtype in (np.complex128, np.float64) else 2e-4
+        assert_allclose(A @ x - b, 0, atol=atol)
+        assert_allclose(x, np.linalg.solve(A, b), atol=5*atol)
+
+    @pytest.mark.parametrize('bdim', [(4,), (4, 3), (2, 3, 4, 3)])
+    @pytest.mark.parametrize('dtype', floating)
+    def test_lstsq(self, bdim, dtype):
+        rng = np.random.default_rng(8342310302941288912051)
+        A = get_random((2, 3, 4, 5), dtype=dtype, rng=rng)
+        b = get_random(bdim, dtype=dtype, rng=rng)
+        res = linalg.lstsq(A, b)
+        x = res[0]
+        if len(bdim) == 1:
+            x = x[..., np.newaxis]
+            b = b[..., np.newaxis]
+        assert_allclose(A @ x - b, 0, atol=2e-6)
+        assert len(res) == 4
+
+    @pytest.mark.parametrize('dtype', floating)
+    def test_clarkson_woodruff_transform(self, dtype):
+        rng = np.random.default_rng(8342310302941288912051)
+        A = get_random((5, 3, 4, 6), dtype=dtype, rng=rng)
+        self.batch_test(linalg.clarkson_woodruff_transform, A,
+                        kwargs=dict(sketch_size=3, rng=311224))
+
+    def test_clarkson_woodruff_transform_sparse(self):
+        rng = np.random.default_rng(8342310302941288912051)
+        A = get_random((5, 3, 4, 6), dtype=np.float64, rng=rng)
+        A = sparse.coo_array(A)
+        message = "Batch support for sparse arrays is not available."
+        with pytest.raises(NotImplementedError, match=message):
+            linalg.clarkson_woodruff_transform(A, sketch_size=3, rng=rng)

scipy/linalg/tests/test_cythonized_array_utils.py

@@ -19,8 +19,6 @@ def test_bandwidth_dtypes():
 def test_bandwidth_non2d_input():
     A = np.array([1, 2, 3])
     raises(ValueError, bandwidth, A)
-    A = np.array([[[1, 2, 3], [4, 5, 6]]])
-    raises(ValueError, bandwidth, A)
 
 
 @pytest.mark.parametrize('T', [x for x in np.typecodes['All']

scipy/linalg/tests/test_decomp.py

@@ -404,6 +404,36 @@ class TestEig:
         assert vr.shape == (0, 0)
         assert vr.dtype == vr_n.dtype
 
+    @pytest.mark.parametrize("include_B", [False, True])
+    @pytest.mark.parametrize("left", [False, True])
+    @pytest.mark.parametrize("right", [False, True])
+    @pytest.mark.parametrize("homogeneous_eigvals", [False, True])
+    @pytest.mark.parametrize("dtype", [np.float32, np.complex128])
+    def test_nd_input(self, include_B, left, right, homogeneous_eigvals, dtype):
+        batch_shape = (3, 2)
+        core_shape = (4, 4)
+        rng = np.random.default_rng(3249823598235)
+        A = rng.random(batch_shape + core_shape).astype(dtype)
+        B = rng.random(batch_shape + core_shape).astype(dtype)
+        kwargs = dict(right=right, homogeneous_eigvals=homogeneous_eigvals)
+
+        if include_B:
+            res = eig(A, b=B, left=left, **kwargs)
+        else:
+            res = eig(A, left=left, **kwargs)
+
+        for i in range(batch_shape[0]):
+            for j in range(batch_shape[1]):
+                if include_B:
+                    ref = eig(A[i, j], b=B[i, j], left=left, **kwargs)
+                else:
+                    ref = eig(A[i, j], left=left, **kwargs)
+
+                if left or right:
+                    for k in range(len(ref)):
+                        assert_allclose(res[k][i, j], ref[k])
+                else:
+                    assert_allclose(res[i, j], ref)
 
 
 class TestEigBanded:
@@ -739,11 +769,11 @@ class TestEigTridiagonal:
     def test_eigvalsh_tridiagonal(self):
         """Compare eigenvalues of eigvalsh_tridiagonal with those of eig."""
         # can't use ?STERF with subselection
-        for driver in ('sterf', 'stev', 'stebz', 'stemr', 'auto'):
+        for driver in ('sterf', 'stev', 'stevd', 'stebz', 'stemr', 'auto'):
             w = eigvalsh_tridiagonal(self.d, self.e, lapack_driver=driver)
             assert_array_almost_equal(sort(w), self.w)
 
-        for driver in ('sterf', 'stev'):
+        for driver in ('sterf', 'stev', 'stevd'):
             assert_raises(ValueError, eigvalsh_tridiagonal, self.d, self.e,
                           lapack_driver=driver, select='i',
                           select_range=(0, 1))
@@ -776,7 +806,7 @@ class TestEigTridiagonal:
         # can't use ?STERF when eigenvectors are requested
         assert_raises(ValueError, eigh_tridiagonal, self.d, self.e,
                       lapack_driver='sterf')
-        for driver in ('stebz', 'stev', 'stemr', 'auto'):
+        for driver in ('stebz', 'stev', 'stevd', 'stemr', 'auto'):
             w, evec = eigh_tridiagonal(self.d, self.e, lapack_driver=driver)
             evec_ = evec[:, argsort(w)]
             assert_array_almost_equal(sort(w), self.w)
@@ -1194,6 +1224,13 @@ def test_svd_gesdd_nofegfault():
         svd(df)
 
 
+def test_gesdd_nan_error_message():
+    A = np.eye(2)
+    A[0, 0] = np.nan
+    with pytest.raises(ValueError, match="NaN"):
+        svd(A, check_finite=False)
+
+
 class TestSVDVals:
 
     @pytest.mark.parametrize('dt', [int, float, np.float32, complex, np.complex64])