code: def flatten_and_batch_shift_indices(indices: torch.Tensor,
                                      sequence_length: int) -> torch.Tensor:
      
      # Shape: (batch_size)
      offsets = get_range_vector(indices.size(0), get_device_of(indices)) * sequence_length
      for _ in range(len(indices.size()) - 1):
          offsets = offsets.unsqueeze(1)
  
      # Shape: (batch_size, d_1, ..., d_n)
      offset_indices = indices + offsets
  
      # Shape: (batch_size * d_1 * ... * d_n)
      offset_indices = offset_indices.view(-1)
      return offset_indices
docs: This is a subroutine for :func:`~batched_index_select`. The given ``indices`` of size
      ``(batch_size, d_1, ..., d_n)`` indexes into dimension 2 of a target tensor, which has size
      ``(batch_size, sequence_length, embedding_size)``. This function returns a vector that
      correctly indexes into the flattened target. The sequence length of the target must be
      provided to compute the appropriate offsets.
  
      .. code-block:: python
  
          indices = torch.ones([2,3], dtype=torch.long)
          # Sequence length of the target tensor.
          sequence_length = 10
          shifted_indices = flatten_and_batch_shift_indices(indices, sequence_length)
          # Indices into the second element in the batch are correctly shifted
          # to take into account that the target tensor will be flattened before
          # the indices are applied.
          assert shifted_indices == [1, 1, 1, 11, 11, 11]
  
      Parameters
      ----------
      indices : ``torch.LongTensor``, required.
      sequence_length : ``int``, required.
          The length of the sequence the indices index into.
          This must be the second dimension of the tensor.
  
      Returns
      -------
      offset_indices : ``torch.LongTensor``