@@ -56,25 +56,6 @@ def _involves_viewed_args(bsym, viewed):
5656 return any (isinstance (p , TensorProxy ) and variableify (p ) in viewed for p in bsym .flat_proxy_args )
5757
5858
59- def _add_output_to_setitem (trace : Trace ) -> Trace :
60- def helper (bsym ):
61- if bsym .sym is not setitem_ :
62- return bsym
63- sym = bsym .sym
64- new_bsym = sym .bind (
65- * bsym .args ,
66- ** bsym .kwargs ,
67- output = bsym .subsymbols [1 ].output ,
68- subsymbols = bsym .subsymbols ,
69- _call_ctx = bsym ._call_ctx ,
70- )
71- return new_bsym
72-
73- new_symbols = [helper (bound_symbol ) for bound_symbol in trace .bound_symbols ]
74- trace .bound_symbols = new_symbols
75- return trace
76-
77-
7859def replace_args_with_alias_map (
7960 computation_trace : Trace ,
8061 alias_tensor_indices : list [list [int ]],
@@ -140,7 +121,6 @@ def insert_alias_updates(computation_trace: Trace, alias_tensor_indices: list[li
140121 # First pass: identify inputs which are views of each other and swap them out with a default,
141122 # reshaping if necessary.
142123 computation_trace , _ = replace_args_with_alias_map (computation_trace , alias_tensor_indices )
143- computation_trace = _add_output_to_setitem (computation_trace )
144124
145125 # Second pass: identify views, their originals, and operands involved in inplace ops
146126 view_groups = []
0 commit comments