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 @@ -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
You can’t perform that action at this time.
0 commit comments