import torch


def rejoin(chunked, initial_shape):
    """
    Rejoins chunked tensor, removing the padding as necessary

    >>> eq = lambda a, b: torch.all(torch.lt(torch.abs(torch.add(a, -b)), 1e-12))
    >>> x = torch.arange(end=4) + 3
    >>> y = torch.arange(end=15) + 2
    >>> mesh = x.view(-1, 1) @ y.view(1, -1)
    >>> mesh = torch.stack([mesh, mesh + 1, mesh + 2], dim=0)

    First we create an array. I don't know why I created it in such a silly way.
    Next, we'll show that chunking/rejoining result in the exact same array,
    despite the fact that some of the chunks are padded!

    >>> mesh.shape
    torch.Size([3, 4, 15])
    >>> chunks(mesh, 3).shape
    torch.Size([1, 2, 5, 3, 3, 3])
    >>> rejoined = rejoin(chunks(mesh, 3), mesh.shape)
    >>> rejoined.shape
    torch.Size([3, 4, 15])
    >>> torch.equal(mesh, rejoined)
    True

    Great! Now we can try specifying a chunk size that is smaller than the
    minimum dimension, and it still works.

    >>> initial = torch.arange(512).view(8, 8, 8)
    >>> chunked = chunks(initial, 9)
    >>> reconstructed = rejoin(chunked, (8, 8, 8))
    >>> torch.equal(initial, reconstructed)
    True

    :param chunked: a chunked tensor created by `chunks`
    :param initial_shape: the initial shape of the tensor before chunking
    :return: tensor in the shape `initial_shape`, dimensions `i` and
        `i + len(initial_shape)` are joined
    """
    indices = []
    padded_shape = []

    for i in range(len(initial_shape)):
        indices.append(i)
        indices.append(i + len(initial_shape))

        padded_shape.append(chunked.shape[i] * chunked.shape[len(initial_shape) + i])

    repermuted = chunked.permute(*indices)
    padded = repermuted.reshape(*padded_shape)

    for i, s in enumerate(initial_shape):
        padded = torch.narrow(padded, i, 0, s)

    return padded