import torch


def reverse_enc(enc, s_lengths):
    """Give the encodings in reverse order along the seq_len dimension,
    reversing only up to the actual sequence length (give by s_lengths) for
    each batch element.

    Inputs:
        enc (tensor) of shape (batch_size, q_seq_len, ...).
        s_lengths (tensor) of shape (batch_size).

    Outputs:
        rev_enc (tensor) of shape (batch_size, q_seq_len, ...).
    """
    seq_len = enc.size(1)
    rev_enc = []
    for enc_b, s_len in zip(enc, s_lengths):
        rev_idx = list(reversed(range(s_len))) + list(range(s_len, seq_len))
        rev_enc.append(enc_b[rev_idx].unsqueeze(0))
    rev_enc = torch.cat(rev_enc, 0)
    return rev_enc