Description
Originally posted by @jessegrabowski in #382 (comment):
Yeah my idea was to cast a wider net. Right now, the rewrite checks for "triangularity" by looking for a Cholesky parent node. So this graph is rewritten:
import pytensor.tensor as pt
import pytensor
A = pt.dmatrix("A")
b = pt.dmatrix('b')
L = pt.linalg.cholesky(A)
x = pt.linalg.solve(L, b)
f = pytensor.function([A, b], x)
But, importantly for PyMC users, this graph is not rewritten:
import pymc as pm
L2, *_ = pm.LKJCholeskyCov.dist(n=3, sd_dist=pm.Exponential.dist(1), eta=1)
x2 = pt.linalg.solve(L2, b)
f2 = pytensor.function([b], x2)
Because the cholesky factorized matrix returned by LKJCholeskyCov isn't actually created via the Cholesky
Op
. My thought was this case could be caught by looking at the new tags instead?
I also wanted to verify that if we have a tagged triangular matrix, pt.linalg.inv(L)
gets re-written to pt.linalg.triangular_solve(L, pt.eye(L.shape[0]))
via sequential application of rewrites. I think it should, but I also didn't check.
Also should we make the outputs of "Cholesky" have such a tag from the get go?
Yes, this was my thinking