Apply rewrites to logp before apply grad

This adds safe rewrites to logp before the grad operator is applied.
This is motivated by #6717, where expensive `cholesky(L.dot(L.T))`
operations are removed. If these remain in the logp graph when the grad
is taken, the resulting dlogp graph contains unnecessary operations.
However this may improve the stability and performance of grad logp
in other situation.

Author: dehorsley

pymc/model.py

@@ -75,6 +75,7 @@
     hessian,
     inputvars,
     replace_rvs_by_values,
+    rewrite_pregrad,
 )
 from pymc.util import (
     UNSET,
@@ -381,6 +382,8 @@ def __init__(
             self._extra_vars_shared[var.name] = shared
             givens.append((var, shared))
 
+        cost = rewrite_pregrad(cost)
+
         if compute_grads:
             grads = pytensor.grad(cost, grad_vars, disconnected_inputs="ignore")
             for grad_wrt, var in zip(grads, grad_vars):
@@ -824,6 +827,7 @@ def dlogp(
                     )
 
         cost = self.logp(jacobian=jacobian)
+        cost = rewrite_pregrad(cost)
         return gradient(cost, value_vars)
 
     def d2logp(
@@ -862,6 +866,7 @@ def d2logp(
                     )
 
         cost = self.logp(jacobian=jacobian)
+        cost = rewrite_pregrad(cost)
         return hessian(cost, value_vars)
 
     @property

pymc/pytensorf.py

@@ -1228,3 +1228,10 @@ def constant_fold(
     return tuple(
         folded_x.data if isinstance(folded_x, Constant) else folded_x for folded_x in folded_xs
     )
+
+
+def rewrite_pregrad(graph):
+    """Apply simplifying or stabilizing rewrites to graph that are safe to use
+    pre-grad.
+    """
+    return rewrite_graph(graph, include=("canonicalize", "stabilize"))