Is there a function to do this operation?

1

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!

2

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.

3

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