-
-
Notifications
You must be signed in to change notification settings - Fork 5.5k
Description
Is your feature request related to a problem? Please describe.
The array API standard and NumPy linalg
functions support batched operation over stacks of matrices, vectors, etc. We occasionally get requests for SciPy to do the same (e.g. gh-16090, gh-17344, gh-20072), and it would help internally sometimes, too. (For instance, I've wanted batch support for many linear algebra operations in multivariate_normal
.) Some of our functions do have this support, so for consistency, we should add support for the rest.
Describe the solution you'd like.
linalg
functions would support N-D input as they do in NumPy.
Conceptually, this means following normal NumPy broadcasting rules with the following adjustment: vectors and matrices, which are slices along the last one or two axes of an array, are treated as though they occupy only one element of the array. We ignore their "core dimensions" for the sake of broadcasting, and these core dimensions are always the last one or two of the array (rather than being specified by an axis
argument).
For implementation, it is more precise to think of the shape of each N-D input as the sum (concatenation) of a batch shape and a core shape (e.g. the shape of a single vector or matrix). The net batch shape is the result of broadcasting the batch shapes of the inputs. The shape of each result array is the sum of the net batch shape and the core shape of each result. We perform the computation by looping over the batch dimensions.
Since solve
and similar functions can accept a 1-D vector or 2-D horizontal stack of column vectors already, there are some special cases, but the descriptions above work in most cases.
Initially, this can be handled throughout linalg
by a decorator (e.g. gh-21462). More performant, lower-level implementations can be added at our leisure and tested against the decorator.
Here's a tracking list. Items marked with +
are included in the array API standard linear algebra extension. (I've marked these because the same decorator could add array API support quite easily.) It's organized in categories that made some sense to me, but it's not perfect. After writing gh-21462, I realized that categorization isn't really needed, anyway; the decorator can be used as-is for almost all cases. I might have left out some functions or neglected to note that they have batch support already. In any case, feel free to modify as you see fit.
Already have batch support (of some kind)
- det+
- norm+
- lu
- cdf2rdf
- solve_circulant
Matrix in | scalar out
- issymmetric (ENH: linalg: add batch support for matrix -> scalar funcs #22127)
- ishermitian (ENH: linalg: add batch support for matrix -> scalar funcs #22127)
- expm_cond (ENH: linalg: add batch support for matrix -> scalar funcs #22127)
Matrices in | vector out
- subspace_angles (ENH: linalg: add batch support to several remaining functions #22147)
- svdvals+ (ENH: linalg: add batch support to several remaining functions #22147)
- eigvals (ENH: linalg: add batch support to remaining eigenvalue functions #22165)
- eigvalsh+ (ENH: linalg: add batch support to several remaining functions #22147)
Matrices in | matrices out
- inv+ (ENH: linalg: add batch support for functions that accept a single array #22133)
- pinv+ (ENH: linalg: add batch support for functions that accept a single array #22133)
- pinvh (ENH: linalg: add batch support for functions that accept a single array #22133)
- null_space (ENH: linalg: add batch support for functions that accept a single array #22133)
- expm (ENH:linalg: expm overhaul and ndarray processing #15079)
- logm (ENH: linalg: add batch support for functions that accept a single array #22133)
- cosm (ENH: linalg: add batch support for functions that accept a single array #22133)
- sinm (ENH: linalg: add batch support for functions that accept a single array #22133)
- tanm (ENH: linalg: add batch support for functions that accept a single array #22133)
- coshm (ENH: linalg: add batch support for functions that accept a single array #22133)
- sinhm (ENH: linalg: add batch support for functions that accept a single array #22133)
- tanhm (ENH: linalg: add batch support for functions that accept a single array #22133)
- signm (ENH: linalg: add batch support for functions that accept a single array #22133)
- sqrtm (ENH: linalg: add batch support for functions that accept a single array #22133)
- funm (ENH: linalg: add batch support for functions that accept a single array #22133)
- fractional_matrix_power+ (ENH: linalg: add batch support for functions that accept a single array #22133)
- matrix_balance (ENH: linalg: add batch support for functions that accept a single array #22133)
- orth (ENH: linalg: add batch support for functions that accept a single array #22133)
- expm_frechet (ENH: linalg: add batch support to several remaining functions #22147)
- orthogonal_procrustes (ENH: linalg: add batch support to several remaining functions #22147)
- khatri_rao (ENH: linalg: add batch support to several remaining functions #22147)
- solve_sylvester (ENH: linalg: add batch support to several remaining functions #22147)
- solve_continous_are (ENH: linalg: add batch support to several remaining functions #22147)
- solve_discrete_are (ENH: linalg: add batch support to several remaining functions #22147)
- solve_continuous_lyapunov (ENH: linalg: add batch support to several remaining functions #22147)
- solve_discrete_lyapunov (ENH: linalg: add batch support to several remaining functions #22147)
Matrices/vectors in | matrices/vectors out
- eig (ENH: linalg.eig: support batched input #21462)
- eigh+ (ENH: linalg: add batch support to several remaining functions #22147)
- svd+ (ENH: linalg: add batch support for functions that accept a single array #22133)
- ldl (ENH: linalg: add batch support for functions that accept a single array #22133)
- cholesky+ (ENH: linalg: add batch support for functions that accept a single array #22133)
- polar (ENH: linalg: add batch support for functions that accept a single array #22133)
- qr+ (ENH: linalg: add batch support for functions that accept a single array #22133)
- rq (ENH: linalg: add batch support for functions that accept a single array #22133)
- qz (ENH: linalg: add batch support to several remaining functions #22147)
- ordqz (ENH: linalg: add batch support to several remaining functions #22147)
- schur (ENH: linalg: add batch support for functions that accept a single array #22133)
- rsf2csf (ENH: linalg: add batch support to several remaining functions #22147)
- hessenberg (ENH: linalg: add batch support for functions that accept a single array #22133)
Vector in | scalar(s) out
Vector(s) in | matrix out
- diagsvd (ENH: linalg: add batch support for functions that accept a single array #22133)
- circulant (ENH:
linalg.circulant
: allow batching #21419) - companion (ENH: vectorize companion matrix function #16090)
- convolution_matrix (ENH: linalg: enable N-D batch support in special matrix functions #21446)
- fiedler (ENH: linalg: enable N-D batch support in special matrix functions #21446)
- fiedler_companion (ENH: linalg: enable N-D batch support in special matrix functions #21446)
- hankel (ENH: linalg: enable N-D batch support in special matrix functions #21446)
- leslie (ENH: linalg: enable N-D batch support in special matrix functions #21446)
Solvers (Matrix and vector(s) in | vector(s) out)
(As written, the _apply_to_batch
decorator won't work with these. One reason is that the RHS can be a vector or array containing multiple vectors, another is that some accept a tuple for the first argument. I'll add support for these special cases when the rest of linalg
is wrapped.)
- solve+ (ENH: linalg: add batch support to linear system solvers #22192)
- lstsq (ENH: linalg: add batch support to linear system solvers #22192)
- solve_triangular (ENH: linalg: add batch support to linear system solvers #22192)
- cho_solve_banded (ENH: linalg: add batch support to linear system solvers #22192)
- cho_solve (ENH: linalg: add batch support to linear system solvers #22192)
- solve_banded (ENH: linalg: add batch support to linear system solvers #22192)
- solveh_banded (ENH: linalg: add batch support to linear system solvers #22192)
- lu_solve (ENH: linalg: add batch support to linear system solvers #22192)
Scalar(s) in | matrix out
(These can't really be batched because it would create a ragged array.)
-
dft -
hadamard -
helmert -
hilbert -
invhilbert -
pascal -
invpascal
Other atypical functions
- cossin - the first argument can be a tuple of arrays (ENH: linalg.cossin: add batch support #22197)
- qr_multiply - core dimensionality of second argument can be 1 or 2 (ENH:
linalg.qr_*
: add batch support #22199) - qr_update - core dimensionality of third/fourth argument can be 1 or 2 (ENH:
linalg.qr_*
: add batch support #22199) - qr_insert - core dimensionality of third argument can be 1 or 2 (ENH:
linalg.qr_*
: add batch support #22199) - qr_delete (ENH:
linalg.qr_*
: add batch support #22199) - clarkson_woodruff_transform - accepts sparse input (ENH: linalg.clarkson_woodruff_transform: add batch support #22154)
- block_diags - accepts an arbitrary number of array arguments (ENH: linalg.block_diag: add batch support #22166)
Didn't categorize yet
- kron (deprecated)
- eig_banded (ENH: linalg: add batch support for functions that accept a single array #22133)
- eigvals_banded (ENH: linalg: add batch support for functions that accept a single array #22133)
- eigvalsh_tridiagonal (ENH: linalg: add batch support to remaining eigenvalue functions #22165)
- eigh_tridiagonal (ENH: linalg: add batch support to remaining eigenvalue functions #22165)
- lu_factor (ENH: linalg: add batch support for functions that accept a single array #22133)
- cholesky_banded (ENH: linalg: add batch support for remaining cholesky functions #22157)
- cho_factor (ENH: linalg: add batch support for remaining cholesky functions #22157)
Need a FutureWarning before batch support can be added:
- toeplitz (ENH: linalg: enable N-D batch support in special matrix functions #21446)
- matmul_toeplitz (DEP: linalg.solve_toeplitz/matmul_toeplitz: warn on n-D
c
,r
#22193) - solve_toeplitz (DEP: linalg.solve_toeplitz/matmul_toeplitz: warn on n-D
c
,r
#22193)
linalg.interpolative
- interp_decomp
- reconstruct_matrix_from_id
- reconstruct_interp_matrix
- reconstruct_skel_matrix
- id_to_svd
- svd
- estimate_spectral_norm
- estimate_spectral_norm_diff
- estimate_rank
Describe alternatives you've considered.
- Requiring users to write their own loops.
- Adding more performant, low-level batch support to functions one at a time