import torch


def collate_fn(batch):
    """Collates a batch of padded variable length inputs"""

    anchor, pos, neg = zip(*batch)

    # get sequence lengths
    anchor_lengths = torch.tensor([x.shape[0] for x in anchor])
    pos_lengths = torch.tensor([x.shape[0] for x in pos])
    neg_lengths = torch.tensor([x.shape[0] for x in neg])

    # padding
    anchor = torch.nn.utils.rnn.pad_sequence(anchor, batch_first=True)
    pos = torch.nn.utils.rnn.pad_sequence(pos, batch_first=True)
    neg = torch.nn.utils.rnn.pad_sequence(neg, batch_first=True)

    return anchor, anchor_lengths, pos, pos_lengths, neg, neg_lengths