def batch_gather(tensors, indices):
  """Gather indices into any dimension rather than dimension 0 of a tensor.

  Gathering is usually performed on dimension 0 of a tensor, such as a word
  embedding matrix. In the case where the target is batched, such as selecting
  a token or a set of tokens for each sequence in the batch, then this function
  will gather the desired elements correctly. For example, given inputs:

  tensor = [[11, 12, 13],
            [21, 22, 23]]
  indices = [[2, 1, 0, 1],
             [0, 1, 1, 2]])

  we expect `batch_gather` to return the following tensor:

  gathered = [[13, 12, 11, 12],
              [21, 22, 22, 23]]

  Args:
    tensors: Tensor or list of Tensors with shape
        [batch_size, sequence_length, d1, d2, ...].
    indices: [batch_size, i1, i2, ...]

  Returns:
    gathered: [batch_size, i1, i2, ..., d1, d2, ...]
  """
  is_singleton = not isinstance(tensors, (list, tuple))

  if is_singleton:
    tensors = [tensors]

  sequence_dim = len(indices.get_shape()) - 1
  if sequence_dim > 1:
    batch_shape = shape(tensors[0])[:sequence_dim]
    batch_size = 1
    for s in batch_shape:
      batch_size *= s
    tensors = [tf.reshape(t, [batch_size] + shape(t)[sequence_dim:])
               for t in tensors]
    indices = tf.reshape(indices, [batch_size] + shape(indices)[sequence_dim:])

  batch_size = shape(indices, 0)
  sequence_length = shape(tensors[0], 1)

  # [batch_size]
  offset = tf.range(batch_size) * sequence_length

  # [batch_size, 1, 1, ...]
  for _ in range(len(shape(indices)) - 1):
    offset = tf.expand_dims(offset, 1)

  offset_indices = indices + offset

  gathered_tensors = []
  for t in tensors:
    # [batch_size * sequence_length, d1, d2, ...]
    flat_tensor = tf.reshape(t, [batch_size * sequence_length] + shape(t)[2:])

    if t.dtype in (tf.int32, tf.int64, tf.bool):
      # Cast to float to enable gathering on GPU.
      flat_tensor = tf.to_float(flat_tensor)

    gathered = tf.gather(flat_tensor, offset_indices)

    if t.dtype in (tf.int32, tf.int64, tf.bool):
      # Undo previous cast.
      gathered = tf.cast(gathered, t.dtype)

    # [batch_size, i1, i2, ..., d1, d2, ...]
    gathered_tensors.append(gathered)

  if sequence_dim > 1:
    gathered_tensors = [tf.reshape(
        t, batch_shape + shape(t)[sequence_dim-1:]) for t in gathered_tensors]
  if is_singleton:
    return gathered_tensors[0]
  else:
    return gathered_tensors
Gather indices into any dimension rather than dimension 0 of a tensor.

  Gathering is usually performed on dimension 0 of a tensor, such as a word
  embedding matrix. In the case where the target is batched, such as selecting
  a token or a set of tokens for each sequence in the batch, then this function
  will gather the desired elements correctly. For example, given inputs:

  tensor = [[11, 12, 13],
            [21, 22, 23]]
  indices = [[2, 1, 0, 1],
             [0, 1, 1, 2]])

  we expect `batch_gather` to return the following tensor:

  gathered = [[13, 12, 11, 12],
              [21, 22, 22, 23]]

  Args:
    tensors: Tensor or list of Tensors with shape
        [batch_size, sequence_length, d1, d2, ...].
    indices: [batch_size, i1, i2, ...]

  Returns:
    gathered: [batch_size, i1, i2, ..., d1, d2, ...]
  
language/labs/exemplar_decoding/utils/tensor_utils.py