Swapping in higher dimention

1

I have a 3-dimentional pytorch array (rows x Column x K) .Lets name the array as ‘X’. I have a list of tuples , ‘L’ . For every tuple (a,b) in ‘L’, I want to swap the values at X[a,b,a] and X[a,b,b]. Is there a fast pytorch operation for this?

(Here rows=Column=K)

2

Hi Kaiser!

You may use pytorch tensor indexing to both access and set the values of
the desired elements, together with a helper tensor to temporarily store the
values being swapped:

Thus:

>>> import torch
>>> print (torch.__version__)
1.13.1
>>>
>>> _ = torch.manual_seed (2023)
>>>
>>> K = 10
>>> N = 5
>>> X = torch.randn (K, K, K)
>>> L = torch.randint (K, (N, 2))
>>>
>>> Xc = X.clone()
>>>
>>> temp = X[L[:, 0], L[:, 1], L[:, 0]]
>>> X[L[:, 0], L[:, 1], L[:, 0]] = X[L[:, 0], L[:, 1], L[:, 1]]
>>> X[L[:, 0], L[:, 1], L[:, 1]] = temp
>>>
>>> L
tensor([[3, 5],
        [7, 6],
        [6, 0],
        [0, 8],
        [5, 9]])
>>> Xc[L[2, 0], L[2, 1]]
tensor([-1.7196, -0.5159, -0.3002,  0.4148, -0.1301, -0.4926,  0.3135, -0.1687,
        -0.5783, -1.2591])
>>> X[L[2, 0], L[2, 1]]
tensor([ 0.3135, -0.5159, -0.3002,  0.4148, -0.1301, -0.4926, -1.7196, -0.1687,
        -0.5783, -1.2591])

Best.

K. Frank

1 Like