Simple solution to roll around first axis:
def roll(x, n):  
    return torch.cat((x[-n:], x[:-n]))

Test like:
x = torch.arange(5)
print("Orig:", x)
print("Roll 2:", roll(x, 2))
print("Roll -2:", roll(x, -2))

Outputs:
Orig: tensor([0, 1, 2, 3, 4])
Roll 2: tensor([3, 4, 0, 1, 2])
Roll -2: tensor([2, 3, 4, 0, 1])

To roll around second axis, use:
def roll_1(x, n):  
    return torch.cat((x[:, -n:], x[:, :-n]), dim=1)

It probably can be generalised, but I didn’t need it.