ENH: linalg: add batch support for functions that accept a single array (scipy#22133)

* ENH: linalg.diagsvd: add batch support

* ENH: linalg.inv: add batch support

* ENH: linalg.null_space: add batch support

* ENH: linalg.sqrtm: add batch support

* ENH: linalg.funm: add batch support

* ENH: linalg.signm: add batch support

* ENH: linalg.fractional_matrix_power: add batch support

* ENH: linalg: add batch support to several functions

* MAINT: linalg.logm: remove warnings filter from test

* ENH: linalg.pinv/pinvh: add batch support

* ENH: linalg.matrix_balance: add batch support

* TST: linalg: test passing arrays by keyword

* ENH: linalg.bandwidth: add batch support

* ENH: linalg.ldl: add batch support

* ENH: linalg.svd: add batch support

* ENH: linalg.cholesky: add batch support

* ENH: linalg.polar: add batch support

* ENH: linalg: add batch support to qr, rq, schur, hessenberg

* MAINT: linalg.clarkson_woodruff_transform: add batch support

* ENH: linalg.orth/lu_factor: add batch support

* ENH: linalg.eig_banded/eigvals_banded: add batch support

* TST: linalg.bandwidth: no more error with 3D input

* MAINT: linalg.cwt: remove batch support to avoid issues with sparse input

Author: mdhaber

scipy/linalg/_basic.py

@@ -8,6 +8,7 @@
 from itertools import product
 import numpy as np
 from numpy import atleast_1d, atleast_2d
+from scipy._lib._util import _apply_over_batch
 from .lapack import get_lapack_funcs, _compute_lwork
 from ._misc import LinAlgError, _datacopied, LinAlgWarning
 from ._decomp import _asarray_validated
@@ -1089,6 +1090,7 @@ def solve_circulant(c, b, singular='raise', tol=None,
 
 
 # matrix inversion
+@_apply_over_batch(('a', 2))
 def inv(a, overwrite_a=False, check_finite=True):
     """
     Compute the inverse of a matrix.
@@ -1499,6 +1501,7 @@ def lstsq(a, b, cond=None, overwrite_a=False, overwrite_b=False,
 lstsq.default_lapack_driver = 'gelsd'
 
 
+@_apply_over_batch(('a', 2))
 def pinv(a, *, atol=None, rtol=None, return_rank=False, check_finite=True):
     """
     Compute the (Moore-Penrose) pseudo-inverse of a matrix.
@@ -1622,6 +1625,7 @@ def pinv(a, *, atol=None, rtol=None, return_rank=False, check_finite=True):
         return B
 
 
+@_apply_over_batch(('a', 2))
 def pinvh(a, atol=None, rtol=None, lower=True, return_rank=False,
           check_finite=True):
     """
@@ -1715,6 +1719,7 @@ def pinvh(a, atol=None, rtol=None, lower=True, return_rank=False,
         return B
 
 
+@_apply_over_batch(('A', 2))
 def matrix_balance(A, permute=True, scale=True, separate=False,
                    overwrite_a=False):
     """

scipy/linalg/_cythonized_array_utils.pyx

@@ -90,6 +90,7 @@ cdef inline void swap_c_and_f_layout(lapack_t *a, lapack_t *b, int r, int c) noe
 # ============================================================================
 
 
+@_apply_over_batch(('a', 2))
 @cython.embedsignature(True)
 def bandwidth(a):
     """Return the lower and upper bandwidth of a 2D numeric array.

scipy/linalg/_decomp.py

@@ -657,6 +657,7 @@ def _check_select(select, select_range, max_ev, max_len):
     return select, vl, vu, il, iu, max_ev
 
 
+@_apply_over_batch(('a_band', 2))
 def eig_banded(a_band, lower=False, eigvals_only=False, overwrite_a_band=False,
                select='a', select_range=None, max_ev=0, check_finite=True):
     """
@@ -1029,6 +1030,7 @@ def eigvalsh(a, b=None, *, lower=True, overwrite_a=False,
                 driver=driver)
 
 
+@_apply_over_batch(('a_band', 2))
 def eigvals_banded(a_band, lower=False, overwrite_a_band=False,
                    select='a', select_range=None, check_finite=True):
     """
@@ -1391,6 +1393,7 @@ def _check_info(info, driver, positive='did not converge (LAPACK info=%d)'):
         raise LinAlgError(("%s " + positive) % (driver, info,))
 
 
+@_apply_over_batch(('a', 2))
 def hessenberg(a, calc_q=False, overwrite_a=False, check_finite=True):
     """
     Compute Hessenberg form of a matrix.

scipy/linalg/_decomp_cholesky.py

@@ -4,6 +4,7 @@
 from numpy import asarray_chkfinite, asarray, atleast_2d, empty_like
 
 # Local imports
+from scipy._lib._util import _apply_over_batch
 from ._misc import LinAlgError, _datacopied
 from .lapack import get_lapack_funcs
 
@@ -43,6 +44,7 @@ def _cholesky(a, lower=False, overwrite_a=False, clean=True,
     return c, lower
 
 
+@_apply_over_batch(('a', 2))
 def cholesky(a, lower=False, overwrite_a=False, check_finite=True):
     """
     Compute the Cholesky decomposition of a matrix.

scipy/linalg/_decomp_ldl.py

@@ -3,13 +3,14 @@
 import numpy as np
 from numpy import (atleast_2d, arange, zeros_like, imag, diag,
                    iscomplexobj, tril, triu, argsort, empty_like)
-from scipy._lib._util import ComplexWarning
+from scipy._lib._util import ComplexWarning, _apply_over_batch
 from ._decomp import _asarray_validated
 from .lapack import get_lapack_funcs, _compute_lwork
 
 __all__ = ['ldl']
 
 
+@_apply_over_batch(('A', 2))
 def ldl(A, lower=True, hermitian=True, overwrite_a=False, check_finite=True):
     """ Computes the LDLt or Bunch-Kaufman factorization of a symmetric/
     hermitian matrix.

scipy/linalg/_decomp_lu.py

@@ -6,6 +6,8 @@
 import numpy as np
 from itertools import product
 
+from scipy._lib._util import _apply_over_batch
+
 # Local imports
 from ._misc import _datacopied, LinAlgWarning
 from .lapack import get_lapack_funcs
@@ -17,6 +19,7 @@
 __all__ = ['lu', 'lu_solve', 'lu_factor']
 
 
+@_apply_over_batch(('a', 2))
 def lu_factor(a, overwrite_a=False, check_finite=True):
     """
     Compute pivoted LU decomposition of a matrix.

scipy/linalg/_decomp_polar.py

@@ -1,10 +1,12 @@
 import numpy as np
+from scipy._lib._util import _apply_over_batch
 from scipy.linalg import svd
 
 
 __all__ = ['polar']
 
 
+@_apply_over_batch(('a', 2))
 def polar(a, side="right"):
     """
     Compute the polar decomposition.

scipy/linalg/_decomp_qr.py

@@ -1,6 +1,8 @@
 """QR decomposition functions."""
 import numpy as np
 
+from scipy._lib._util import _apply_over_batch
+
 # Local imports
 from .lapack import get_lapack_funcs
 from ._misc import _datacopied
@@ -23,6 +25,7 @@ def safecall(f, name, *args, **kwargs):
     return ret[:-2]
 
 
+@_apply_over_batch(('a', 2))
 def qr(a, overwrite_a=False, lwork=None, mode='full', pivoting=False,
        check_finite=True):
     """
@@ -366,6 +369,7 @@ def qr_multiply(a, c, mode='right', pivoting=False, conjugate=False,
     return (cQ,) + raw[1:]
 
 
+@_apply_over_batch(('a', 2))
 def rq(a, overwrite_a=False, lwork=None, mode='full', check_finite=True):
     """
     Compute RQ decomposition of a matrix.

scipy/linalg/_decomp_schur.py

@@ -3,7 +3,7 @@
 from numpy import asarray_chkfinite, single, asarray, array
 from numpy.linalg import norm
 
-
+from scipy._lib._util import _apply_over_batch
 # Local imports.
 from ._misc import LinAlgError, _datacopied
 from .lapack import get_lapack_funcs
@@ -14,6 +14,7 @@
 _double_precision = ['i', 'l', 'd']
 
 
+@_apply_over_batch(('a', 2))
 def schur(a, output='real', lwork=None, overwrite_a=False, sort=None,
           check_finite=True):
     """

scipy/linalg/_decomp_svd.py

@@ -2,14 +2,18 @@
 import numpy as np
 from numpy import zeros, r_, diag, dot, arccos, arcsin, where, clip
 
+from scipy._lib._util import _apply_over_batch
+
 # Local imports.
 from ._misc import LinAlgError, _datacopied
 from .lapack import get_lapack_funcs, _compute_lwork
 from ._decomp import _asarray_validated
 
+
 __all__ = ['svd', 'svdvals', 'diagsvd', 'orth', 'subspace_angles', 'null_space']
 
 
+@_apply_over_batch(('a', 2))
 def svd(a, full_matrices=True, compute_uv=True, overwrite_a=False,
         check_finite=True, lapack_driver='gesdd'):
     """
@@ -249,6 +253,7 @@ def svdvals(a, overwrite_a=False, check_finite=True):
                check_finite=check_finite)
 
 
+@_apply_over_batch(('s', 1))
 def diagsvd(s, M, N):
     """
     Construct the sigma matrix in SVD from singular values and size M, N.
@@ -301,6 +306,7 @@ def diagsvd(s, M, N):
 
 # Orthonormal decomposition
 
+@_apply_over_batch(('A', 2))
 def orth(A, rcond=None):
     """
     Construct an orthonormal basis for the range of A using SVD
@@ -349,6 +355,7 @@ def orth(A, rcond=None):
     return Q
 
 
+@_apply_over_batch(('A', 2))
 def null_space(A, rcond=None, *, overwrite_a=False, check_finite=True,
                lapack_driver='gesdd'):
     """