Is there a function to do this operation?

When I have a tensor “x” (size: T), I want to get a tensor “dist” (size: (D, T)) like the following. (D < T)

dist = torch.tensor(
        [x[0], x[1], x[2], ..., x[T-1]],
        [x[-1], x[0], x[1], ...,x[T-2]],
        [x[-2], x[-1], x[0], ...,x[T-3]],
        [x[-D], x[1-D], x[2-D], ..., x[T-D]]

What is the best way to do this?
Thank you!

This will do the job:

import torch
D = 3
T = 4
x = torch.rand(T)

idx = torch.arange(T)
idx = (idx[None,:] - idx[:, None])%T

x.repeat(D).index_select(0, idx.view(-1)[:D * T]).view(D, T)

If T is very big, then you could have some memory issues, and a for loop would be better for constructing the idx tensor.

I’m not sure it is the best way, but it should be good enough for many contexts.

Thanks for the answer!
That’s a very clever approach. I would never have thought of that.

I’m going to do this on a tensor of size batch_size×T with batch_size=20,T=6000.

1 Like