Pytorch : Pass from grid to image coordinate convention for use of grid_sample
I need to interpolate some deformation grid in PyTorch and decided to use the function grid_sample (see doc). I need to reshape grids as images back and forth with theses conventions :

N Batch size
D grid depth (for 3D images)
H grid height
W grid width
d grid dimension (=2 or 3)

images format is (N,2,H,W) in 2D (resp. (N,3,D,H,W) in 3D)
when grid format is (N,H,W,2) (resp. (N,D,H,W,3)in 3D)
I can't use reshape or view because they do not arrange the data as I wish. I need to have (for example)
grid_in_grid_convention[0,:,:,0] == grid_in_image_convention[0,0,:,:]

I came up with theses functions to make the reshaping which works fine but I am sure that there is a more compact/fast way to do so. What are your taught ?
def grid2im(grid):
    &quot;&quot;&quot;Reshape a grid tensor into an image tensor
        2D  [T,H,W,2] -&gt; [T,2,H,W]
        3D  [T,D,H,W,2] -&gt; [T,D,H,W,3]
    &quot;&quot;&quot;
    if grid.shape[0] == 1 and grid.shape[-1] == 2: # 2D case, batch =1
        return torch.stack((grid[0,:,:,0],grid[0,:,:,1]),dim = 0).unsqueeze(0)

    elif grid.shape[0] == 1 and grid.shape[-1] == 3: # 3D case, batch =1 
        return torch.stack((grid[0,:,:,:,0],grid[0,:,:,:,1],grid[0,:,:,:,2]),
                            dim = 0).unsqueeze(0)
    
    elif grid.shape[-1] == 2:
        N,H,W,d = grid.shape
        temp = torch.zeros((N,H,W,d))
        for n in range(N):
            temp[n,:,:,:] = torch.stack((grid[n,:,:,0],grid[n,:,:,1]),dim = 0).unsqueeze(0)
        return temp
    
    elif grid.shape[-1] == 3:
        N,D,H,W,d =grid.shape
        temp = torch.zeros((N,D,H,W,d))
        for n in range(N):
            temp[n,:,:,:,:] = torch.stack((grid[n,:,:,:,0],
                                           grid[n,:,:,:,1],
                                           grid[n,:,:,:,2]),
                            dim = 0).unsqueeze(0)
    else:
        raise ValueError(&quot;input argument expected is [N,H,W,2] or [N,D,H,W,3]&quot;,
                         &quot;got &quot;+str(grid.shape)+&quot; instead.&quot;)

def im2grid(image):
    &quot;&quot;&quot;Reshape an image tensor into a grid tensor
        2D case [T,2,H,W]   -&gt;  [T,H,W,2]
        3D case [T,3,D,H,W] -&gt;  [T,D,H,W,3]
    &quot;&quot;&quot;
    # No batch 
    if image.shape[0:2] == (1,2):
        return torch.stack((image[0,0,:,:],image[0,1,:,:]),dim= 2).unsqueeze(0)
    elif image.shape[0:2] == (1,3):
        return torch.stack((image[0,0,:,:],image[0,1,:,:],image[0,2,:,:]),
                           dim = 2).unsqueeze(0)
    # Batch size &gt; 1
    elif image.shape[0] &gt; 0 and image.shape[1] == 2 :
        N,d,H,W = image.shape
        temp = torch.zeros((N,H,W,d))
        for n in range(N):
            temp[n,:,:,:] = torch.stack((image[n,0,:,:],image[n,1,:,:]),dim= 2).unsqueeze(0)
        return temp
    elif image.shape[0] &gt; 0 and image.shape[1] == 3 :
        N,d,D,H,W = image.shape
        temp = torch.zeros((N,D,H,W,d))
        for n in range(N):
            temp[n,:,:,:] = torch.stack((image[n,0,:,:],
                                         image[n,1,:,:],
                                         image[n,2,:,:]),
                           dim = 2).unsqueeze(0)
        return temp
    else:
        raise ValueError(&quot;input argument expected is [1,2,H,W] or [1,3,D,H,W]&quot;,
                         &quot;got &quot;+str(image.shape)+&quot; instead.&quot;)


You can use transpose(). For example, if you want [T,H,W,2] -&gt; [T,2,H,W] for a tensor, say grid, you can do
grid = grid.transpose(2,3).transpose(1,2)

