Skip to content

Commit 45a33ad

Browse files
committed
Copy stack trace in InplaceGraphRewriter
1 parent 1ebd078 commit 45a33ad

File tree

1 file changed

+5
-1
lines changed

1 file changed

+5
-1
lines changed

pytensor/tensor/rewriting/elemwise.py

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -199,12 +199,14 @@ def apply(self, fgraph):
199199
prof["nb_eager_inconsistent"] += 1
200200
else:
201201
prof["nb_replaced"] += 1
202+
copy_stack_trace(node.outputs, inplace_node.outputs)
202203
continue
203204

204205
# If it fails or doesn't match the desired inplace pattern, try one output/input at a time
205206
tried_inputs = set()
206207
inplace_pattern = {}
207208
replaced = False
209+
original_node = node
208210
for (o, _), (i, _) in sorted_candidate_pairs:
209211
if o not in inplace_pattern and i not in tried_inputs:
210212
inplace_pattern[o] = [i]
@@ -225,7 +227,9 @@ def apply(self, fgraph):
225227
prof["nb_inconsistent"] += 1
226228
# The input, not the output caused inconsistencies
227229
inplace_pattern.pop(o)
228-
prof["nb_replaced"] += replaced
230+
if replaced:
231+
copy_stack_trace(original_node.outputs, node.outputs)
232+
prof["nb_replaced"] += replaced
229233

230234
return prof
231235

0 commit comments

Comments
 (0)