Add lower triangular tags to allow chol rewrites

Since pymc-devs/pytensor#303, `cholesky(L.dot(L.T))` will rewritten to L
if `L.tag.lower_triangular=True`. This change adds these where
appropriate.

Fixes #6717.

Author: dehorsley

pymc/distributions/multivariate.py

@@ -132,6 +132,9 @@ def quaddist_matrix(cov=None, chol=None, tau=None, lower=True, *args, **kwargs):
         chol = pt.as_tensor_variable(chol)
         if chol.ndim != 2:
             raise ValueError("chol must be two dimensional.")
+
+        # tag as lower triangular to enable pytensor rewrites of chol(l.l') -> l
+        chol.tag.lower_triangular = True
         cov = chol.dot(chol.T)
 
     return cov

pymc/math.py

@@ -443,11 +443,17 @@ def expand_packed_triangular(n, packed, lower=True, diagonal_only=False):
     elif lower:
         out = pt.zeros((n, n), dtype=pytensor.config.floatX)
         idxs = np.tril_indices(n)
-        return pt.set_subtensor(out[idxs], packed)
+        # tag as lower triangular to enable pytensor rewrites
+        out = pt.set_subtensor(out[idxs], packed)
+        out.tag.lower_triangular = True
+        return out
     elif not lower:
         out = pt.zeros((n, n), dtype=pytensor.config.floatX)
         idxs = np.triu_indices(n)
-        return pt.set_subtensor(out[idxs], packed)
+        # tag as upper triangular to enable pytensor rewrites
+        out = pt.set_subtensor(out[idxs], packed)
+        out.tag.upper_triangular = True
+        return out
 
 
 class BatchedDiag(Op):

tests/test_pytensorf.py

@@ -30,6 +30,7 @@
 from pytensor.tensor.random.basic import normal, uniform
 from pytensor.tensor.random.op import RandomVariable
 from pytensor.tensor.random.var import RandomStateSharedVariable
+from pytensor.tensor.slinalg import Cholesky
 from pytensor.tensor.subtensor import AdvancedIncSubtensor, AdvancedIncSubtensor1
 from pytensor.tensor.var import TensorVariable
 
@@ -878,3 +879,34 @@ def replacement_fn(var, replacements):
         [new_x], _ = _replace_vars_in_graphs([x], replacement_fn=replacement_fn)
 
         assert new_x.eval() > 50
+
+
+def test_mvnormal_no_cholesky_op():
+    """
+    Test MvNormal likelihood when using Cholesky factor parameterization does not unnecessarily
+    recompute the cholesky factorization
+    Reversion test of #6717
+    """
+    with pm.Model() as m:
+        n = 3
+        sd_dist = pm.HalfNormal.dist(shape=n)
+        chol, corr, sigmas = pm.LKJCholeskyCov("cov", n=n, eta=1, sd_dist=sd_dist)
+        mu = np.zeros(n)
+        data = np.ones((10, n))
+        pm.MvNormal("y", mu=mu, chol=chol, observed=data)
+
+    contains_cholesky_op = lambda fgraph: any(
+        isinstance(node.op, Cholesky) for node in fgraph.apply_nodes
+    )
+
+    logp = m.compile_logp()
+    assert not contains_cholesky_op(logp.f.maker.fgraph)
+
+    dlogp = m.compile_dlogp()
+    assert not contains_cholesky_op(dlogp.f.maker.fgraph)
+
+    d2logp = m.compile_d2logp()
+    assert not contains_cholesky_op(d2logp.f.maker.fgraph)
+
+    logp_dlogp = m.logp_dlogp_function()
+    assert not contains_cholesky_op(logp_dlogp._pytensor_function.maker.fgraph)