File tree Expand file tree Collapse file tree 1 file changed +5
-1
lines changed
pytensor/tensor/rewriting Expand file tree Collapse file tree 1 file changed +5
-1
lines changed Original file line number Diff line number Diff line change @@ -199,12 +199,14 @@ def apply(self, fgraph):
199
199
prof ["nb_eager_inconsistent" ] += 1
200
200
else :
201
201
prof ["nb_replaced" ] += 1
202
+ copy_stack_trace (node .outputs , inplace_node .outputs )
202
203
continue
203
204
204
205
# If it fails or doesn't match the desired inplace pattern, try one output/input at a time
205
206
tried_inputs = set ()
206
207
inplace_pattern = {}
207
208
replaced = False
209
+ original_node = node
208
210
for (o , _ ), (i , _ ) in sorted_candidate_pairs :
209
211
if o not in inplace_pattern and i not in tried_inputs :
210
212
inplace_pattern [o ] = [i ]
@@ -225,7 +227,9 @@ def apply(self, fgraph):
225
227
prof ["nb_inconsistent" ] += 1
226
228
# The input, not the output caused inconsistencies
227
229
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
229
233
230
234
return prof
231
235
You can’t perform that action at this time.
0 commit comments