import torch


def impute(x, lengths, channels_last=True):
    """ Set elements of a batch of a sequence of tensors to zero according to sequence lengths.
    :param x: A tensor with shape [batch, time_step, ...] or [batch, ..., time_step]
    :param lengths: A tensor with shape [batch]
    :param channels_last: A bool. If true, the time_step dimension is the second dimension, otherwise the last.

    :returns: A tensor with the same shape of x, with elements time_step > corresponding length set to 0.
    """

    if channels_last:
        max_length = x.shape[1]
    else:
        max_length = x.shape[-1]
    mask = torch.arange(max_length, device=lengths.device)[None, :] < lengths[:, None]  # [B, T]
    for _ in range(len(x.shape) - 2):
        if channels_last:
            mask = mask.unsqueeze(-1)
        else:
            mask = mask.unsqueeze(1)
    return x * mask