[torch.jit.script] Support tensor indexing after an Ellipsis

[alanhdu]

🚀 Feature

Support Tensor indexing after an ellipsis: right now the following code:

@torch.jit.script
def zero_diag(inputs: torch.Tensor) -> torch.Tensor:
    dim = inputs.shape[-1]
    inputs[..., torch.range(0, dim), torch.range(0, dim)] = 0

fails with:

Ellipses followed by tensor indexing is currently not supported:
  File "<ipython-input-105-deac139d15e2>", line 4
def zero_diag(inputs: torch.Tensor) -> torch.Tensor:
    dim = inputs.shape[-1]
    inputs[..., torch.range(0, dim), torch.range(0, dim)] = 0
    ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE

Motivation

For our application, we would like to manipulate the diagonal of a batch of matrices, which is most easily expressed by manipulating inputs[..., range(N), range(N)]. Unfortunately, Torch does not currently support this when using torch.jit.script.

Another use-case is to vectorize a batch of symmetric matrices: ideally, we could do something like:

ind = torch.triu_indices(N, N)
vectors = matrices[..., ind[0], ind[1]]

cc @suo

[ailzhang]

Thanks @alanhdu for the feature request! @driazati IIRC you have a list of language features we'll be adding, does this fall into our plan?

[vfdev-5]

Various padding mode of tensor can be also implemented with tensor indexing. And it would be very nice to have jit support for that too.

[Abhijit-2592]

Any update on this?