ENH: linalg: add batch support to remaining eigenvalue functions (scipy#22165)

* ENH: linalg.eigh_tridiagonal: add batch support

* ENH: linalg.eigvalsh_tridiagonal: add batch support

* ENH: linalg.eigvals: add batch support

* TST: linalg.cdf2rdf: add test of batch support

* Revert "TST: linalg.cdf2rdf: add test of batch support"

This reverts commit 40bdaf3.

* TST: linalg: generalize name of TestBatch

Author: mdhaber

scipy/linalg/_decomp.py

@@ -838,6 +838,7 @@ def eig_banded(a_band, lower=False, eigvals_only=False, overwrite_a_band=False,
     return w, v
 
 
+@_apply_over_batch(('a', 2), ('b', 2))
 def eigvals(a, b=None, overwrite_a=False, check_finite=True,
             homogeneous_eigvals=False):
     """
@@ -1126,6 +1127,7 @@ def eigvals_banded(a_band, lower=False, overwrite_a_band=False,
                       select_range=select_range, check_finite=check_finite)
 
 
+@_apply_over_batch(('d', 1), ('e', 1))
 def eigvalsh_tridiagonal(d, e, select='a', select_range=None,
                          check_finite=True, tol=0., lapack_driver='auto'):
     """
@@ -1207,6 +1209,7 @@ def eigvalsh_tridiagonal(d, e, select='a', select_range=None,
         check_finite=check_finite, tol=tol, lapack_driver=lapack_driver)
 
 
+@_apply_over_batch(('d', 1), ('e', 1))
 def eigh_tridiagonal(d, e, eigvals_only=False, select='a', select_range=None,
                      check_finite=True, tol=0., lapack_driver='auto'):
     """

scipy/linalg/tests/test_batch.py

@@ -26,10 +26,11 @@ def get_nearly_hermitian(shape, dtype, atol, rng):
     return A + At + noise
 
 
-class TestOneArrayIn:
-    # Test the functions that accept one array argument
+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):
+    def batch_test(self, fun, arrays, *, core_dim=2, n_out=1, kwargs=None, dtype=None,
+                   broadcast=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
@@ -53,7 +54,8 @@ def batch_test(self, fun, arrays, core_dim=2, n_out=1, kwargs=None, dtype=None):
         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.
-        arrays = np.broadcast_arrays(*arrays)
+        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]):
@@ -210,7 +212,7 @@ def test_eigvals_banded(self, dtype, rng):
 
     @pytest.mark.parametrize('two_in', [False, True])
     @pytest.mark.parametrize('fun_n_nout', [(linalg.eigh, 1), (linalg.eigh, 2),
-                                            (linalg.eigvalsh, 1)])
+                                            (linalg.eigvalsh, 1), (linalg.eigvals, 1)])
     @pytest.mark.parametrize('dtype', floating)
     def test_eigh(self, two_in, fun_n_nout, dtype, rng):
         fun, n_out = fun_n_nout
@@ -291,6 +293,15 @@ def test_rsf2cs(self, dtype, rng):
     @pytest.mark.parametrize('dtype', floating)
     def test_cholesky_banded(self, dtype, rng):
         ab = get_random((5, 4, 3, 6), dtype=dtype, rng=rng)
-        ab[..., 0, 0] = 0
         ab[..., -1, :] = 10  # make diagonal dominant
         self.batch_test(linalg.cholesky_banded, ab)
+
+    @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):
+        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)