support aten::index.Tensor

[ruoqianguo]

Description

aten::index.Tensor(Tensor self, Tensor?[] indices) -> (Tensor) is a complicated PyTorch function which can receive several indices. To focus on Video Swin Transformer, this version converter only supports one index. I will improve this converter in a future release.

Fixes #766

Type of change

Please delete options that are not relevant and/or add your own.

• New feature (non-breaking change which adds functionality)

Checklist:

• My code follows the style guidelines of this project (You can use the linters)
• I have performed a self-review of my own code
• I have commented my code, particularly in hard-to-understand areas and hacks
• I have made corresponding changes to the documentation
• I have added tests to verify my fix or my feature
• New and existing unit tests pass locally with my changes

[narendasan]

Can you clarify what this means? Can I only do something like tensor[0] or something?

[ruoqianguo]

In Pytorch, 'aten::index.Tensor' can receive several indices, like this:

x = torch.rand(5, 7, 10)
idx0 = torch.tensor([0, 0, 0, 0])
idx1 = torch.tensor([1, 1, 1, 1])
idx2 = torch.tensor([2, 2, 2, 2])
y1 = x[idx0, idx1, idx2] # result is like tensor([x[0, 1, 2], x[0, 1, 2], x[0, 1, 2], x[0, 1, 2]])
y2 = x[idx0] # result is like tensor([x[0, :, :], x[0, :, :], x[0, :, :], x[0, :, :]])

I only support one index as input currently like y2 = x[idx0]. We can only process tensors[0].

[ruoqianguo]

In Pytorch, 'aten::index.Tensor' can receive several indices, like this:

x = torch.rand(5, 7, 10)
idx0 = torch.tensor([0, 0, 0, 0])
idx1 = torch.tensor([1, 1, 1, 1])
idx2 = torch.tensor([2, 2, 2, 2])
y1 = x[idx0, idx1, idx2] # result is like tensor([x[0, 1, 2], x[0, 1, 2], x[0, 1, 2], x[0, 1, 2]])
y2 = x[idx0] # result is like tensor([x[0, :, :], x[0, :, :], x[0, :, :], x[0, :, :]])

I only support one index as input currently like y2 = x[idx0]. We can only process tensors[0].

Hi @narendasan Is this explanation clear?

[narendasan]

Ok can we improve this message to convey what this means

[ruoqianguo]

Hi @narendasan i have added some annotation and improved this message.

[narendasan]

LGTM

[ncomly-nvidia]

Fixes #468