Need for Improvement:How to get every index of a tensor

Suppose we have [3,2,2] matrix
So, I need all index like:
[[0,0,0],[0,0,1]…[2,1,1]]

If I understand correctly, you want something like this.

If not, please let me know.

Hope this helps :blush:

import torch
import itertools

t = torch.rand((3, 2, 2))
shape = t.shape

y = torch.tensor(list(itertools.product(range(len(shape)), repeat=len(shape))))

for i, max in enumerate(shape):
    y = y[ y[:, i] < max ]

print(y)
# Output:
tensor([[0, 0, 0],
        [0, 0, 1],
        [0, 1, 0],
        [0, 1, 1],
        [1, 0, 0],
        [1, 0, 1],
        [1, 1, 0],
        [1, 1, 1],
        [2, 0, 0],
        [2, 0, 1],
        [2, 1, 0],
        [2, 1, 1]])
1 Like

Yes, that’s what I need, you helped a lot.
:hugs:
If this process could be simpler, that’s the best.
Thank you

1 Like

If what you want is to access every value consecutively, then you can do something like this:

t = torch.rand((3, 2, 2))

print(t)

for i in range(len(t.view(-1))):
    print(t.view(-1)[i])

Yeah, I know that.
Unfortunately, I can’t access these values directly.
Anyway, thanks and wish you all be good.

1 Like

@Matias_Vasquez
Hi, I find that there might have been a bug existing,

shape = [64,3,7,7]
idx = torch.tensor(list(itertools.product(range(len(shape)), repeat=len(shape))))
for i,right_b in enumerate(shape):
    idx = idx[ idx[:,i]<right_b]

Oops, inside the product it has to be the maximum number, not the length

shape = [64,3,7,7]
idx = torch.tensor(list(itertools.product(range(max(shape)), repeat=len(shape))))
for i,right_b in enumerate(shape):
    idx = idx[ idx[:,i]<right_b]

This might be more efficient, since the other answer has to create all of the different permutations, taking the maximum in every dimension.


t = torch.rand((64, 3, 7, 7))

B, C, H, W = t.shape

y = torch.tensor([[b, c, h, w] for b in range(B) for c in range(C) for h in range(H) for w in range(W)])

print(y)

Yes, but dimensions of tensor might be changed, so I need some adaptive ways to solve that.

Here should be a generic solution for n dimensions.

Not the most elegant, but also works.

import torch
import numpy as np


t = torch.rand((64, 3, 7, 7))
z = torch.tensor([[0] * len(t.shape)] * np.prod(t.shape))

shape = t.shape

for i in range(-1, -1-len(shape), -1):
    if i == -1:
        z[:, i] = torch.tensor(
            [j for j in range(shape[i])] * (int)(np.prod(shape)/(np.prod(shape[i:]))) 
        )
    else:
        z[:, i] = torch.tensor(
            [j for j in range(shape[i]) for _ in range(np.prod(shape[i+1:]))] * (int)(np.prod(shape)/(np.prod(shape[i:]))) 
        )

print(z)