def batch_gather(params, indices):
  """Gather a batch of indices from a batch of params."""
  batch_size = tf.shape(indices)[0]
  num_idx = tf.shape(indices)[1]
  brange = tf.cast(tf.range(batch_size), indices.dtype)
  bindices = tf.tile(tf.expand_dims(brange, 1), (1, num_idx))
  gather_indices = tf.stack([bindices, indices], axis=2)
  return tf.gather_nd(params, gather_indices)
Gather a batch of indices from a batch of params.
language/labs/drkit/model_fns.py