code: def gather_2d(params, indices):
      """2D version of tf.gather.
  
      This is a batched version of tf.gather(), i.e. it applies tf.gather() to
      each batch separately.
      Example:
        params = [[10, 11, 12, 13, 14],
                  [20, 21, 22, 23, 24]]
        indices = [[0, 0, 1, 1, 1, 2],
                   [1, 3, 0, 0, 2, 2]]
        result = [[10, 10, 11, 11, 11, 12],
                  [21, 23, 20, 20, 22, 22]]
      This method is copied from
        https://github.com/fstahlberg/tensor2tensor-usr/blob/master/usr/utils.py
      which is published under Apache 2.
  
      Args:
        params: A [batch_size, n, ...] tensor with data
        indices: A [batch_size, num_indices] int32 tensor with indices into params.
          Entries must be smaller than n
  
      Returns:
        The result of tf.gather() on each entry of the batch.
      """
      batch_size = tf.shape(params)[0]
      num_indices = tf.shape(indices)[1]
      batch_indices = tf.tile(
          tf.expand_dims(tf.range(batch_size), 1), [1, num_indices]
      )
      # batch_indices is [[0,0,0,0,...],[1,1,1,1,...],...]
      gather_nd_indices = tf.stack([batch_indices, indices], axis=2)
      return tf.gather_nd(params, gather_nd_indices)
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import logging
  from tensor2tensor.layers import common_attention
  from tensor2tensor.layers import common_layers
  from tensor2tensor.layers import transformer_layers
  from tensor2tensor.models import transformer
  from tensor2tensor.utils import registry
  from tensor2tensor.utils import t2t_model
  import tensorflow.compat.v1
label_desc: Use image node_id 5 for calling a global function with example usage: gather_2d(params, indices) and returns: tf