Rewrite local_inplace_cholesky to use make_inplace helper function

Author: jessegrabowski

pytensor/tensor/rewriting/basic.py

@@ -135,6 +135,18 @@ def alloc_like(
     return rval
 
 
+def make_inplace(node, inplace_prop="inplace"):
+    op = getattr(node.op, "core_op", node.op)
+    props = op._props_dict()
+    if props[inplace_prop]:
+        return False
+
+    props[inplace_prop] = True
+    inplace_op = type(op)(**props)
+
+    return inplace_op.make_node(*node.inputs).outputs
+
+
 def register_useless(
     node_rewriter: Union[RewriteDatabase, NodeRewriter, str], *tags, **kwargs
 ):

pytensor/tensor/rewriting/linalg.py

@@ -10,6 +10,7 @@
 from pytensor.tensor.math import Dot, Prod, _matrix_matrix_matmul, log, prod
 from pytensor.tensor.nlinalg import MatrixInverse, det
 from pytensor.tensor.rewriting.basic import (
+    make_inplace,
     register_canonicalize,
     register_specialize,
     register_stabilize,
@@ -313,15 +314,9 @@ def local_log_prod_sqr(fgraph, node):
         # returns the sign of the prod multiplication.
 
 
-cholesky_no_inplace = Cholesky(overwrite_a=False)
-cholesky_inplace = Cholesky(overwrite_a=True)
-
-
-@node_rewriter([cholesky_no_inplace], inplace=True)
+@node_rewriter([Cholesky], inplace=True)
 def local_inplace_cholesky(fgraph, node):
-    new_out = [cholesky_inplace(*node.inputs)]
-    copy_stack_trace(node.outputs, new_out)
-    return new_out
+    return make_inplace(node, "overwrite_a")
 
 
 # After destroyhandler(49.5) but before we try to make elemwise things

pytensor/tensor/slinalg.py

@@ -47,12 +47,12 @@ class Cholesky(Op):
     # TODO: for specific dtypes
     # TODO: LAPACK wrapper with in-place behavior, for solve also
 
-    __props__ = ("lower", "destructive", "on_error")
+    __props__ = ("lower", "overwrite_a", "on_error")
     gufunc_signature = "(m,m)->(m,m)"
 
     def __init__(self, *, lower=True, on_error="raise", overwrite_a=False):
         self.lower = lower
-        self.destructive = overwrite_a
+        self.overwrite_a = overwrite_a
         if on_error not in ("raise", "nan"):
             raise ValueError('on_error must be one of "raise" or ""nan"')
         self.on_error = on_error
@@ -72,15 +72,15 @@ def perform(self, node, inputs, outputs):
         (z,) = outputs
 
         try:
-            if x.flags["C_CONTIGUOUS"]:
+            if x.flags["C_CONTIGUOUS"] and self.overwrite_a:
                 # Inputs to the LAPACK functions need to be exactly as expected for overwrite_a to work correctly,
                 # see https://github.com/scipy/scipy/issues/8155#issuecomment-343996798
                 x = scipy.linalg.cholesky(
-                    x.T, lower=not self.lower, overwrite_a=self.destructive
+                    x.T, lower=not self.lower, overwrite_a=self.overwrite_a
                 ).T
             else:
                 x = scipy.linalg.cholesky(
-                    x, lower=self.lower, overwrite_a=self.destructive
+                    x, lower=self.lower, overwrite_a=self.overwrite_a
                 )
 
         except scipy.linalg.LinAlgError:

tests/tensor/rewriting/test_linalg.py

@@ -313,12 +313,12 @@ def test_local_inplace_cholesky():
     L = cholesky(X, overwrite_a=False, lower=True)
     f = function([pytensor.In(X, mutable=True)], L)
 
-    assert not L.owner.op.core_op.destructive
+    assert not L.owner.op.core_op.overwrite_a
 
     nodes = f.maker.fgraph.toposort()
     for node in nodes:
         if isinstance(node, Cholesky):
-            assert node.destructive
+            assert node.overwrite_a
             break
 
     X_val = np.random.normal(size=(10, 10)).astype(config.floatX)