Don't fallback for aten.index when there's a None in the middle of the indexing.

[Chillee]

🚀 The feature, motivation and pitch

import torch
torch.set_default_device('cuda')

torch.manual_seed(420)

x = torch.arange(3*4*5).view(3, 4, 5)

def f(x):
    return x[torch.tensor([1, 2]), :, torch.tensor([2, 3])]

print(f(x))
print(torch.compile(f)(x))
>>>
tensor([[22, 27, 32, 37],
        [43, 48, 53, 58]], device='cuda:0')

This gets lowered as a fallback kernel due to this check (https://github.com/pytorch/pytorch/blob/main/torch/_inductor/lowering.py#L2567). Let's handle that case correctly.

The semantics of having a None there are identical to having a torch.arange there and broadcasting the rest of the indices appropriately.

def f(x):
    return x[torch.tensor([1, 2]).unsqueeze(-1), torch.arange(4), torch.tensor([2, 3]).unsqueeze(-1)]

Alternatives

No response

Additional context

No response

cc @ezyang @msaroufim @wconstab @bdhirsh @anijain2305 @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @Xia-Weiwen @wenzhe-nrv @jiayisunx @peterbell10 @ipiszy @yf225 @chenyang78 @kadeng @muchulee8 @aakhundov @ColinPeppler