Update base for Update on "[2/N] Dynamic Shape: Enable dynamic shape support for aoti_eager"

cc voznesenskym penguinwu jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx peterbell10 ipiszy yf225 chenyang78 kadeng muchulee8 ColinPeppler amjames desertfire chauhang

[ghstack-poisoned]
This commit is contained in:
Wang, Eikan
2024-06-06 14:40:24 +00:00
parent 929afcf437
commit be047b923f

View File

@ -203,15 +203,15 @@ def _unlift_graph(mod, gm, graph_signature):
user_input_mutations = graph_signature.user_inputs_to_mutate
output_tokens = graph_signature.output_tokens
for idx, out in enumerate(outputs):
value = None
if idx < len(buffer_mutations) + len(user_input_mutations) + len(output_tokens):
if out.name in graph_signature.buffers_to_mutate:
mutated_outputs.append(graph_signature.buffers_to_mutate[out.name])
elif out.name in graph_signature.user_inputs_to_mutate:
mutated_outputs.append(graph_signature.user_inputs_to_mutate[out.name])
else:
mutated_outputs.append(None)
else:
mutated_outputs.append(None)
if out.name in buffer_mutations:
value = buffer_mutations[out.name]
elif out.name in user_input_mutations:
value = user_input_mutations[out.name]
mutated_outputs.append(value)
unlifted_gm = _unlift(
gm,