Wrong graph optimization with dynamic shape #2577
Description
import torch
import torch.nn as nn
class Model(nn.Module):
def forward(self, x):
x[:2] = 0
return x
def main():
model = Model()
args = (torch.rand(10),)
dim = torch.export.Dim("dim")
dynamic_shapes = {"x": {0: dim}}
torch.onnx.export(
model,
args,
"model_test.onnx",
dynamic_shapes=dynamic_shapes,
dynamo=True,
)
if __name__ == '__main__':
main()The output graph will be optimized to:
Graph with optimize=False will be:
And This is what happen: collapse_slice2 found that input shape and output shape of Slice (the right one in the graph) are "same" (they are both [None]), so it is optmized to Identity. Then it match the rule ScatterAllDynamic, the whole graph is optimized to an Identity op.
The root cause of trouble is ir.SymbolicDim allow None dim, and treat None == None. A simple solution to fix it is convert all None dim to a unique string, just need to add these two lines in SymbolicDim.__init__:
if value is None:
value = "_s" + str(id(self))