#20153

I believe you need 2 passes for this. Take this example

@torch.jit.script
def f():
    x = torch.ones(10, 9, 8, 7, 6)
    return x[..., None, None].shape

which results in [10, 9, 8, 7, 6, 1, 1]
vs

@torch.jit.script
def f():
    x = torch.ones(10, 9, 8, 7, 6)
    return x[..., None, None, :].shape

which results in [10, 9, 8, 7, 1, 1, 6]
After only processing x[..., None, None we don't know whether we should be creating a new dimension at the end of the dimension list or somewhere in the middle. What we do depends on the elements to the right of it.

Thus, I do 2 passes - one to collect all the dimensions that the index operations operate on, and another that executes the index operations.

This still doesn't work for an ellipse index followed by a tensor index, but it wasn't working previously either.