import torch


def _permute_tensor(input, permutation):
    """
    Permute elements of tensor given a matrix of permutation indices

    Args:
        input: torch.Tensor to be permuted
        permuation: Permutation of the indeces of all elements of input (flat)

    Returns:
        Permuted input
    """
    # Cache the original dimensions
    dimensions = input.size()

    # Apply the permutation to the flattened tensor
    output_flat = torch.index_select(input.view(-1), 0, permutation)

    # Restore original dimensions
    output = output_flat.view(dimensions)

    return output