ENH: linalg: add batch support for remaining cholesky functions (scipy#22157)

* ENH: linalg.cho_factor: add batch support

* ENH: linalg.cholesky_banded: add batch support

Author: j-bowhay

scipy/linalg/_decomp_cholesky.py

@@ -105,6 +105,7 @@ def cholesky(a, lower=False, overwrite_a=False, check_finite=True):
     return c
 
 
+@_apply_over_batch(("a", 2))
 def cho_factor(a, lower=False, overwrite_a=False, check_finite=True):
     """
     Compute the Cholesky decomposition of a matrix, to use in cho_solve
@@ -246,6 +247,7 @@ def cho_solve(c_and_lower, b, overwrite_b=False, check_finite=True):
     return x
 
 
+@_apply_over_batch(("ab", 2))
 def cholesky_banded(ab, overwrite_ab=False, lower=False, check_finite=True):
     """
     Cholesky decompose a banded Hermitian positive-definite matrix

scipy/linalg/tests/test_batch.py

@@ -160,7 +160,8 @@ def test_bandwidth(self, dtype, 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)])
+    @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):
         fun, n_out = fun_n_out
@@ -286,3 +287,10 @@ def test_rsf2cs(self, dtype, rng):
         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):
+        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)