Get a 10x10 patch from a 100x100 pytorch tensor with torus style wrap around the boundries
How can I get a 10x10 patch from a 100x100 pytorch tensor, with the added constraint that if a patch would go outside the boundaries of the array, then it wraps around the edges (as if the array was a torus, with the top joined to the bottom, and the left joined to the right)?
I wrote this code that will do the job, I'm looking for something more elegant, efficient and clear:
def shift_matrix(a, distances) -&gt; Tensor:
  x, y = distances
  a = torch.cat((a[x:], a[0:x]), dim=0)
  a = torch.cat((a[:, y:], a[:, :y]), dim=1)
  return a

def randomly_shift_matrix(a) -&gt; Tensor:
  return shift_matrix(a, np.random.randint(low = 0, high = a.size()))

def random_patch(a, size) -&gt; Tensor:
  full_shifted_matrix = randomly_shift_matrix(a)
  return full_shifted_matrix[0:size[0], 0:size[1]]

I feel like something with negative index slices should work.  I haven't found it though.
You can see the code in google colab here.

You are looking for torch.roll
def random_patch(a, size) -&gt; Tensor:
  shifts = np.random.randint(low = 0, high = a.size())
  return torch.roll(a, shifts=shifts, dims=(0, 1))[:size[0], :size[1]]

