search

rgl-epfl / cholespy

An easily integrable Cholesky solver on CPU and GPU
BSD 3-Clause "New" or "Revised" License
210 stars 18 forks source link

Batched implementation #26

Open jacarvalho opened 1 year ago

jacarvalho commented 1 year ago

Hey, First, thanks a lot for the repo! Is there a way to use cholespy with a batch of matrices? If not, what would you suggest as the best approach? Thanks!

bathal1 commented 1 year ago

Hi,

Cholespy supports batches of vectors, but for a fixed matrix. The reason is that the amount and structure of computation required to solve the system varies depending on the sparsity pattern of the matrix, so this makes it a poor fit for parallel execution on the GPU.

If you have batches of matrices, chances are that you are not using a sparse representation, so this package is probably not the right one for your application.

jacarvalho commented 1 year ago

I have a batch of sparse matrices that, unfortunately, are not stored as sparse matrices since pytorch still has very poor support of operations for these kinds of matrices. (It would be a lot of work to restructure the code at this point, just for a marginal gain.) One thing I thought of was to construct a block diagonal of the batch of sparse matrices, which would increase memory but hopefully would make it faster to solve with your method than using torch.cholesky_solve.