def batched_index(values, indices, keepdims=None):
  """Equivalent to `values[:, indices]`.

  Performs indexing on batches and sequence-batches by reducing over
  zero-masked values. Compared to indexing with `tf.gather` this approach is
  more general and TPU-friendly, but may be less efficient if `num_values`
  is large. It works with tensors whose shapes are unspecified or
  partially-specified, but this op will only do shape checking on shape
  information available at graph construction time. When complete shape
  information is absent, certain shape incompatibilities may not be detected at
  runtime! See `indexing_ops_test` for detailed examples.

  Args:
    values: tensor of shape `[B, num_values]` or `[T, B, num_values]`
    indices: tensor of shape `[B]` or `[T, B]` containing indices.
    keepdims: If `True`, the returned tensor will have an added 1 dimension at
      the end (e.g. `[B, 1]` or `[T, B, 1]`).

  Returns:
    Tensor of shape `[B]` or `[T, B]` containing values for the given indices.

  Raises: ValueError if values and indices have sizes that are known
    statically (i.e. during graph construction), and those sizes are not
    compatible (see shape descriptions in Args list above).
  """
  with tf.name_scope("batch_indexing", values=[values, indices]):
    values = tf.convert_to_tensor(values)
    indices = tf.convert_to_tensor(indices)
    assert_compatible_shapes(values.shape, indices.shape)

    one_hot_indices = tf.one_hot(
        indices, tf.shape(values)[-1], dtype=values.dtype)
    return tf.reduce_sum(values * one_hot_indices, axis=-1, keepdims=keepdims)
Equivalent to `values[:, indices]`.

  Performs indexing on batches and sequence-batches by reducing over
  zero-masked values. Compared to indexing with `tf.gather` this approach is
  more general and TPU-friendly, but may be less efficient if `num_values`
  is large. It works with tensors whose shapes are unspecified or
  partially-specified, but this op will only do shape checking on shape
  information available at graph construction time. When complete shape
  information is absent, certain shape incompatibilities may not be detected at
  runtime! See `indexing_ops_test` for detailed examples.

  Args:
    values: tensor of shape `[B, num_values]` or `[T, B, num_values]`
    indices: tensor of shape `[B]` or `[T, B]` containing indices.
    keepdims: If `True`, the returned tensor will have an added 1 dimension at
      the end (e.g. `[B, 1]` or `[T, B, 1]`).

  Returns:
    Tensor of shape `[B]` or `[T, B]` containing values for the given indices.

  Raises: ValueError if values and indices have sizes that are known
    statically (i.e. during graph construction), and those sizes are not
    compatible (see shape descriptions in Args list above).
  
trfl/indexing_ops.py