[FIXED] Parallel Cholesky decomposition in PyTorch GPU

Issue

To get a differentiable term containing the determinant of a D-dimensional positive-definite matrix C (differential entropy of a multivariate Gaussian in my case), I can use:

torch.log2(torch.potrf(C).diag()).sum() + D / 2.0 * (np.log2(2 * np.pi * np.e))

potrf(C) performs Cholesky decomposition, whose diagonal elements' log values sum to the log determinant divided by 2.

I want to call potrf on a mini-batch of matrices, so that calling potrf on a Tensor of shape (N, D, D) produces N different Cholesky decompositions.

At the moment I can only call potrf() repeatedly in a Python loop, which is a poor use of the GPU's parallel compute capability and as a result runs about 3 times slower than CPU.

Is it possible to launch Cholesky decomposition on GPU in parallel with PyTorch?

Solution

Batch Cholesky decomposition is now available in PyTorch.

Along with batch inverse(), etc.

---

---

For older versions of Pytorch

You are looking for Batch Cholesky decomposition. It's not implemented in Pytorch at the moment, but there is an open issue and plan to add it in the future.

I only know of a Batch LU factorization available in Pytorch 0.4. You could use it to get something similar:

det(D) = det(P)det(L)det(U)

where the determinant of P is (-1)^t and the determinants of P and U are the product of the diagonal elements.

Answered By - iacolippo