Support for batched `cupyx.scipy.sparse.linalg.spsolve`

Description

Currently, cupy.linalg.solve supports batch processing. Such functionality for cupyx.scipy.sparse.linalg.spsolve would be a great asset and would prevent slow, manual loops.

For example: N being the batch dimension A x = b , with

• A: Sparse matrix with dimension (N, M, M) or as list with N-Sparse matrices (M, M).

• b: Dense matrix with dimension (N, M) or as list with N-Dense vectors (M).

• x: Dense matrix with dimension (N, M) or as list with N-Dense vectors (M).

Additional Information

No response

[asorie]

In the meantime, I found another solution to enable “batched” spsolve. I combined all my A input matrices into one global matrix and reshaped the b vectors into one vector. After that, instead of many systems of equations, you can solve just the one. Then the output vector x can be reshaped into the desired batched shape. The basic idea came from this post.

[leofang]

@kmaehashi @emcastillo can we reopen this for now? batched spsolve is one of the most frequent asks I’ve run into and it’d be nice to discuss further.