Skip to content

Commit 1438953

Browse files
committed
Copy stack trace in InplaceGraphRewriter
1 parent 591c8d4 commit 1438953

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
@@ -205,12 +205,14 @@ def apply(self, fgraph):
205205
prof["nb_eager_inconsistent"] += 1
206206
else:
207207
prof["nb_replaced"] += 1
208+
copy_stack_trace(node.outputs, inplace_node.outputs)
208209
continue
209210

210211
# If it fails or doesn't match the desired inplace pattern, try one output/input at a time
211212
tried_inputs = set()
212213
inplace_pattern = {}
213214
replaced = False
215+
original_node = node
214216
for (o, _), (i, _) in sorted_candidate_pairs:
215217
if o not in inplace_pattern and i not in tried_inputs:
216218
inplace_pattern[o] = [i]
@@ -233,7 +235,9 @@ def apply(self, fgraph):
233235
prof["nb_inconsistent"] += 1
234236
# The input, not the output caused inconsistencies
235237
inplace_pattern.pop(o)
236-
prof["nb_replaced"] += replaced
238+
if replaced:
239+
copy_stack_trace(original_node.outputs, node.outputs)
240+
prof["nb_replaced"] += replaced
237241

238242
return prof
239243

0 commit comments

Comments
 (0)