Skip to content

Expand linalg rewrites based on "triangular" tag #385

Open
@ricardoV94

Description

@ricardoV94

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

Metadata

Metadata

Assignees

No one assigned

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions