Keep stack trace in random_make_inplace

Author: ricardoV94

pytensor/tensor/random/rewriting/basic.py

@@ -51,7 +51,10 @@ def random_make_inplace(fgraph, node):
         props = op._props_dict()
         props["inplace"] = True
         new_op = type(op)(**props)
-        return new_op.make_node(*node.inputs).outputs
+        new_outputs = new_op.make_node(*node.inputs).outputs
+        for old_out, new_out in zip(node.outputs, new_outputs):
+            copy_stack_trace(old_out, new_out)
+        return new_outputs
 
     return False
 

tests/tensor/random/rewriting/test_basic.py

@@ -7,7 +7,7 @@
 from pytensor.compile.mode import Mode
 from pytensor.graph.basic import Constant
 from pytensor.graph.fg import FunctionGraph
-from pytensor.graph.rewriting.basic import EquilibriumGraphRewriter
+from pytensor.graph.rewriting.basic import EquilibriumGraphRewriter, check_stack_trace
 from pytensor.graph.rewriting.db import RewriteDatabaseQuery
 from pytensor.tensor import constant
 from pytensor.tensor.elemwise import DimShuffle
@@ -96,6 +96,8 @@ def test_inplace_rewrites():
     )
     assert np.array_equal(new_out.owner.inputs[1].data, [])
 
+    assert check_stack_trace(f)
+
 
 def test_inplace_rewrites_extra_props():
     class Test(RandomVariable):