How to index/slice the last dimension of a PyTorch tensor/numpy array of unknown dimensions

Question

For example, if I have a 2D tensor X, I can do slicing X[:,1:]; if I have a 3D tensor Y, then I can do similar slicing for the last dimension like Y[:,:,1:].

What is the right way to do the slicing when given a tensor Z of unknown dimension? How about a numpy array?

Thanks!

Answer 1

PyTorch support NumPy-like indexing so you can use Ellipsis(...)

>>> z[..., -1:]

Example:

>>> x                     # (2,2) tensor
tensor([[0.5385, 0.9280],
        [0.8937, 0.0423]])
>>> x[..., -1:]
tensor([[0.9280],
        [0.0423]])

>>> y                     # (2,2,2) tensor
tensor([[[0.5610, 0.8542],
         [0.2902, 0.2388]],

        [[0.2440, 0.1063],
         [0.7201, 0.1010]]])
>>> y[..., -1:]
tensor([[[0.8542],
         [0.2388]],

        [[0.1063],
         [0.1010]]])

• Ellipsis (...) expands to the number of : objects needed for the selection tuple to index all dimensions. In most cases, this means that length of the expanded selection tuple is x.ndim. There may only be a single ellipsis present.