code: 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
          )
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import tensorflow.compat.v1
label_desc: Use image node_id 2 for calling a global function with example usage: batched_index(values, indices, keepdims) and returns: tf