Torch - Randomly choose dimension of tensor to slice

Question

I have a torch tensor of size (1 x n x n x n) and I would like to randomly choose one of the last 3 dimensions to randomly slice at s and then do. For example it could output the below tensors with equal probability;

(1 x s x n x n)
(1 x n x s x n)
(1 x n x n x s)

I realise I could just do a few if else statements but I am curious if there is a "neater" option using a function like torch.random(1,4) to select the dimension.

Answer 1

assuming that you want to narrow a block of s elements randomly, out of n elements.

Let's use :narrow.

n = 100
s = 20
x = torch.randn(1, n, n, n)
y = x:narrow(torch.random(2, 4), torch.random(1, n - s + 1), s)