import torch


def gather_positions(sequence, positions):
    """Gathers the vectors at the specific positions over a minibatch.

  Args:
    sequence: A [batch_size, seq_length] or
        [batch_size, seq_length, depth] tensor of values
    positions: A [batch_size, n_positions] tensor of indices

  Returns: A [batch_size, n_positions] or
    [batch_size, n_positions, depth] tensor of the values at the indices
  """
    shape = sequence.size()
    depth_dimension = (len(shape) == 3)
    if depth_dimension:
        B, L, D = shape
    else:
        B, L = shape
        D = 1

    position_shift = torch.unsqueeze(L * torch.arange(B, device=positions.device), -1)
    flat_positions = torch.reshape(positions + position_shift, [-1])
    flat_sequence = torch.reshape(sequence, [B * L, D])
    gathered = flat_sequence[flat_positions]
    if depth_dimension:
        return torch.reshape(gathered, [B, -1, D])
    else:
        return torch.reshape(gathered, [B, -1])