Interleaving a set of channels during concatenation?
I am trying to perform a quaternion space concatenation which requires the four dimensions r,i,j,k to be concatenated. According to quaternion theory, we cannot apply the torch.cat function directly as they would mess up the components as r channels have to be concatenated with r channels and so on. I managed to perform this action using the quaternion_concat function adapted from here. However, when I use this for networks like densenet, this takes a long time due to multiple for loops concatenations.
Code: Let me give an example, tensor_1 and tensor_2 are two tensors that need to be concatenated, and it has 16 channels each. That means it has 4 channels of r,i,j,k respectively stacked together. I used the torch.chunk function to separate these and concatenate them separately and finally combine them back. Is there any way I could efficiently perform this?
import torch

def quarternion_concat(x, dim=2):
    output = [[] for i in range(4)]
    for _x in x:
        sp = torch.chunk(_x, 4, dim=dim)
        for i in range(4):
            output[i].append(sp[i])

    final = []
    for o in output:
        o = torch.cat(o, dim)
        final.append(o)

    return torch.cat(final, dim)


tensor_1 = torch.randn((1, 16, 64, 64), requires_grad=False)
tensor_2 = torch.randn((1, 16, 64, 64), requires_grad=False)

tensor3 = quarternion_concat([tensor_1, tensor_2], dim=1)


Based on the previous answer, I came up with a solution that utilizes class to compute the indices and use index_select to rearrange the concatenation. The only drawback is that the number of inputs and channels are required before execution.
import torch
from time import time


def quarternion_concat(x, dim=2):
    output = [[] for i in range(4)]
    for _x in x:
        sp = torch.chunk(_x, 4, dim=dim)
        for i in range(4):
            output[i].append(sp[i])

    final = []
    for o in output:
        o = torch.cat(o, dim)
        final.append(o)

    return torch.cat(final, dim)


tensor_1 = torch.randn((1, 16, 64, 64), requires_grad=False)
tensor_2 = torch.randn((1, 16, 64, 64), requires_grad=False)
tensor_3 = torch.randn((1, 16, 64, 64), requires_grad=False)
tensor_4 = torch.randn((1, 16, 256, 256), requires_grad=False)
tensor_5 = torch.randn((1, 16, 256, 256), requires_grad=False)
tensor_6 = torch.randn((1, 16, 256, 256), requires_grad=False)
tensor_7 = torch.randn((1, 16, 1024, 1024), requires_grad=False)
tensor_8 = torch.randn((1, 16, 1024, 1024), requires_grad=False)
tensor_9 = torch.randn((1, 16, 1024, 1024), requires_grad=False)


def quarternion_concat2(x, dim=1):
    output = torch.empty(tuple([i if I != dim else i * len(x) for I, i in enumerate(x[0].shape)])).cuda()
    # output = torch.cat(x, dim=dim)
    inds = torch.arange(output.shape[dim]).view([-1, 4 * len(x)]).cuda()
    for i in range(len(x)):
        output[:, inds[:, i * 4:i * 4 + 4].flatten()] = x[i]
    return output


class concat(torch.nn.Module):
    def __init__(self, no_of_inputs, total_number_of_channel):
        super().__init__()
        temp = torch.chunk(torch.arange(total_number_of_channel).view([-1, 4, 4]), no_of_inputs, dim=0)
        self.register_buffer('indx', torch.cat(temp, dim=2).flatten())

    def forward(self, x):
        output = torch.cat(x, dim=1).index_select(dim=1, index=self.indx)
        return output


def time_funcs(x, str, N=100):
    x = [i.cuda() for i in x]
    print(str)
    s = time()
    for _ in range(N):
        tensor4 = quarternion_concat(x, dim=1)
    print('First method: {:1.7f}'.format((time() - s) / 1000))

    s = time()
    for _ in range(N):
        tensor5 = quarternion_concat2(x, dim=1)
    print('Second method: {:1.7f}'.format((time() - s) / 1000))

    m = concat(len(x), 16 * len(x)).cuda()
    s = time()
    for _ in range(N):
        tensor6 = m(x)
    print('Third method: {:1.7f}'.format((time() - s) / 1000))


time_funcs([tensor_1, tensor_2], 'Two small')
time_funcs([tensor_1, tensor_2, tensor_3], 'Three small')
time_funcs([tensor_4, tensor_5], 'Two big')
time_funcs([tensor_4, tensor_5, tensor_6], 'Three big')
time_funcs([tensor_7, tensor_8], 'Two huge')
time_funcs([tensor_7, tensor_8, tensor_9], 'Three huge')



Below are the results when run on a GPU.
Two small
First method: 0.0000054
Second method: 0.0000203
Third method: 0.0000031
Three small
First method: 0.0000081
Second method: 0.0000295
Third method: 0.0000022
Two big
First method: 0.0000050
Second method: 0.0001637
Third method: 0.0000028
Three big
First method: 0.0000056
Second method: 0.0002477
Third method: 0.0000022
Two huge
First method: 0.0000051
Second method: 0.0025740
Third method: 0.0000028
Three huge
First method: 0.0000062
Second method: 0.0038824
Third method: 0.0000028

