Extracting subtensors from tensors

1

I’m trying to index a tensor t to compute norms of subtensors included in t.
Suppose that t is of shape (8, 3, 6, 6), and that I have an indexing tensor idx of shape (12, 3). idx contains the coordinates of 12 elements I want to extract from t for each element in the batch. I then compute the norm of the subtensor:

I’ve tried t[idx], t[:, idx] which both don’t work. I have one solution which works :

idx_ = (...,) + tuple(idx.T)
norm = torch.linalg.norm(t[idx_], dim=-1)

But when profiling this, forming idx_ is slower than computing the norm, and it seems quite unwieldy and unnatural. Would there be any faster and more natural way to perform this indexing?

2

Could you explain how the idx tensor should be used in indexing t?
Since dim1 of idx is of size 3, would you like to apply each of these values to dim1 to dim3 in t?

3

Yes, exactly. The 3 values are the indexes in dim1 to dim3.
The resulting tensor should be of shape (8, 12) give or take some squeezing.

4

Would another representation for idx make this work?

5

If I understand the use case correctly, this should work:

x = torch.randn(8, 3, 6, 6)
idx = torch.cat((torch.randint(0, 3, (12, 1)), torch.randint(0, 6, (12, 2))), 1)

res = x[:, idx[:, 0], idx[:, 1], idx[:, 2]]
1 Like