"arbitrary number of leading dimensions" only supports up to 4 dimensions

📚 The doc issue

I was attempting do an affine transform of medical imaging data organized using two channel dimensions (SLICE and SEQUENCE) for a total of five: BATCH x SLICE x SEQUENCE x HEIGHT x WIDTH, but torchvision.transforms.functional.affine only supports up to four dimensions.

import torchvision
x = torch.empty((1, 320, 320))
x_t = torchvision.transforms.functional.affine(x, angle=0, translate=(0,0), scale=1, shear=0)
print(x.shape, x_t.shape)

x = torch.empty((1, 1, 320, 320))
x_t = torchvision.transforms.functional.affine(x, angle=0, translate=(0,0), scale=1, shear=0)
print(x.shape, x_t.shape)

x = torch.empty((1, 1, 1, 320, 320))
x_t = torchvision.transforms.functional.affine(x, angle=0, translate=(0,0), scale=1, shear=0)
print(x.shape, x_t.shape)

x = torch.empty((1, 1, 1, 1, 320, 320))
x_t = torchvision.transforms.functional.affine(x, angle=0, translate=(0,0), scale=1, shear=0)
print(x.shape, x_t.shape)

Resulting output and error message:

torch.Size([1, 320, 320]) torch.Size([1, 320, 320])
torch.Size([1, 1, 320, 320]) torch.Size([1, 1, 320, 320])
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-211-2ee8b4ff09ce> in <module>
      9 
     10 x = torch.empty((1, 1, 1, 320, 320))
---> 11 x_t = torchvision.transforms.functional.affine(x, angle=0, translate=(0,0), scale=1, shear=0)
     12 print(x.shape, x_t.shape)
     13 

~/bmelseg/venv/lib/python3.6/site-packages/torchvision/transforms/functional.py in affine(img, angle, translate, scale, shear, interpolation, fill, resample, fillcolor)
   1123     translate_f = [1.0 * t for t in translate]
   1124     matrix = _get_inverse_affine_matrix([0.0, 0.0], angle, translate_f, scale, shear)
-> 1125     return F_t.affine(img, matrix=matrix, interpolation=interpolation.value, fill=fill)
   1126 
   1127 

~/bmelseg/venv/lib/python3.6/site-packages/torchvision/transforms/functional_tensor.py in affine(img, matrix, interpolation, fill)
    696     # grid will be generated on the same device as theta and img
    697     grid = _gen_affine_grid(theta, w=shape[-1], h=shape[-2], ow=shape[-1], oh=shape[-2])
--> 698     return _apply_grid_transform(img, grid, interpolation, fill=fill)
    699 
    700 

~/bmelseg/venv/lib/python3.6/site-packages/torchvision/transforms/functional_tensor.py in _apply_grid_transform(img, grid, mode, fill)
    645         img = torch.cat((img, dummy), dim=1)
    646 
--> 647     img = grid_sample(img, grid, mode=mode, padding_mode="zeros", align_corners=False)
    648 
    649     # Fill with required color

~/bmelseg/venv/lib/python3.6/site-packages/torch/nn/functional.py in grid_sample(input, grid, mode, padding_mode, align_corners)
   4009         align_corners = False
   4010 
-> 4011     return torch.grid_sampler(input, grid, mode_enum, padding_mode_enum, align_corners)
   4012 
   4013 

RuntimeError: grid_sampler(): expected 4D or 5D input and grid with same number of dimensions, but got input with sizes [1, 1, 1, 320, 320] and grid with sizes [1, 320, 320, 2]

cc @vfdev-5 @datumbox

Read torchvision/transforms/functional.py and functional_tensor.py, starting at affine and _apply_grid_transform, then run the five-dimensional reproduction from the issue. Done means affine accepts the shown higher-dimensional input and preserves the expected shape without the grid_sampler dimensionality error.