# Tensor advanced slicing/ indexing

I am working on a vision problem that requires shift in preprocessing.

img is 3d tensor shaped [N, m, m].
shiftmap is an int tensor containing starting index, shaped [N, 2].
K is a constant.

I did the shift in a for-loop:

``````import torch

ret = torch.zeros(N, K, K)
for i in range(N):
xx = shiftmap[i, 0]
yy = shiftmap[i , 1]
ret[i, ...] = img[i, xx: xx+K, yy: yy+K]

``````

Is there an advanced Tensor slicing operation that can do the shift in O(1) instead of a for-loop?

Thanks a ton!

Hi Ultraman!

Yes. One scheme would be to create index-tensor â€śtemplatesâ€ť for the `xx`
and `yy` indices of `img` and index into those templates with `shiftmap` to get
the actual `xx` and `yy` indices. Then index into `img` itself with tensor indexing
(together with a rather trivial tensor index for the batch (`N`) dimension) to get
the final shifted result.

Here is an illustration:

``````>>> import torch
>>> torch.__version__
'1.13.0'
>>>
>>> _ = torch.manual_seed (2023)
>>>
>>> N = 5
>>> m = 10
>>>
>>> K = 3
>>>
>>> L = m - K + 1   # limit of allowed shiftmap (plus 1)
>>>
>>> img = torch.randn (N, m, m)            # sample img
>>> shiftmap = torch.randint (L, (N, 2))   # sample shiftmap
>>>
>>> # with for-loop
>>> ret = torch.zeros(N, K, K)
>>> for i in range(N):
...     xx = shiftmap[i, 0]
...     yy = shiftmap[i , 1]
...     ret[i, ...] = img[i, xx: xx+K, yy: yy+K]
...
>>> # loop-free version
>>> ind0 = torch.arange (N)[:, None, None].repeat (1, K, K)
>>> ind2 = torch.arange (m).repeat (m + 1).reshape (m, m + 1)[0:L, 0:K].unsqueeze (1).expand (L, K, K)
>>> ind1 = ind2.transpose (1, 2)
>>> retB = img[ind0, ind1[shiftmap[:, 0]], ind2[shiftmap[:, 1]]]
>>>
>>> torch.equal (retB, ret)
True
``````

Technically speaking, you canâ€™t do this in O(1) time (even with a loop-free
algorithm). No matter how you slice it you have to create and populate your
result tensor of shape `[N, K, K]`, so the best your algorithm can be is
O (N * K**2). (Of course, pushing those operations out of python loops down
into pytorch tensor operations will offer a huge speed-up.)

Best.

K. Frank

1 Like