I just realized that this problem is even more complicate. If we simply rename the second T_reshape to T_reshape_1, then you still cannot match the TE compute unless we rename the corresponding one in the compute DAG as well. However, printing the compute DAG and schedule are different functions and it would be unsafe to maintain a such mapping. While the meta schedule that @junrushao et all are working on has resolved this issue, I would say the best way for now is probably just to manually match them, as we know that the line with ... = tuple(name.op.axis) + tuple(name.op.reduce_axis) is always the starting point of scheduling a stage.
1 Like