Generalize `linalg.cholesky` beyond 2D arrays

[ricardoV94]

import numpy as np
import aesara.tensor as at

x = np.full((2, 3, 3), np.eye(3))
np.linalg.cholesky(x)  # broadcast operation fine
at.linalg.cholesky(x)  # AssertionError x.ndim == 2

[Sayam753]

Hi
I have started to look into this issue to make progress on pymc-devs/pymc#5424. The source code of Cholesky Op defines a L_Op method:

aesara/aesara/tensor/slinalg.py

Line 70 in da86d35

def L_op(self, inputs, outputs, gradients):

To generalize linalg.cholesky beyond 2d arrays, do we need to update this method as well alongside changing perform method?

[brandonwillard]

To generalize linalg.cholesky beyond 2d arrays, do we need to update this method as well alongside changing perform method?

Yes, as well as Op.infer_shape and everything else. It's always possible to raise exceptions in those cases, but that would make the Op/new capability much less useful.

[lucianopaz]

@Sayam753, did you start working on this? I'd like to try doing it if not.

[Sayam753]

@lucianopaz no, feel free to work on the issue. I am still wrapping my head around how the grad operations in Op work.

[lucianopaz]

I've got a working Cholesky Op, but I'm kind of struggling with the gradient. Aesara's scipy wrapper around Cholesky uses solve_upper_triangular, which expects 2-D arrays. Scipy, doesn't seem to have anything that we can use to solve a linear equation with a batch of triangular matrices. Numpy, has solve that should be slower and maybe less stable than an algorithm that exploits the triangular nature of the matrix, but it can handle batches of matrices. The thing is that aesara does not have an Op that wraps numpy solve. Should I open an issue to add one? I feel it shouldn't be done on a PR that is trying to get cholesky to work.

As a side note, there are also a bunch of functions tril, diag and diagonal that expected matrices, but it was easy to code up extensions that could accommodate batches of matrices simply using advanced indexing. Should those be made available in nlinalg module as well as cholesky?

[ricardoV94]

@lucianopaz The perform method of the Op is free to do whatever it wants, including Python loops across batches, or just call a np.vectorized version of the scipy solve if we have reason to believe that one is better.