import torch


def pack(x: torch.Tensor, sequence_lengths: torch.Tensor):
    """
    Packs `x` such that it combines its batch (0) and time (1) axis and removes
    any padded values in between. It can be reverted with `unpack`.

    .. note::

        This is different from `pack_padded_sequence` in that it does not
        interleave the time steps and does not return a `PackedSequence`.
    """
    assert len(sequence_lengths) == len(x)
    return torch.cat([x_[:l] for x_, l in zip(x, sequence_lengths)])