Skip to content

ENH: linalg: N-D array (batch) support  #21466

@mdhaber

Description

@mdhaber

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

Matrices in | vector out

Matrices in | matrices out

Matrices/vectors in | matrices/vectors out

Vector in | scalar(s) out

Vector(s) in | matrix out

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.)

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

Didn't categorize yet

Need a FutureWarning before batch support can be added:

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

Metadata

Metadata

Assignees

No one assigned

    Labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions