def get_batch_emb_idxs(emb_idx_list, max_num_batch_emb, config):
  """Get unique embedding idxs within the batch."""
  batch_size = tf.shape(emb_idx_list)[0]
  flatten_shape = (batch_size * config["max_num_nts"] *
                   config["max_num_numerator_nodes"],)
  # <int>[num_unquie_emb_idxs], <int>[batch_size * num_nts * num_nodes]
  batch_unique_emb_idxs, batch_emb_idx_list = tf.unique(
      tf.reshape(emb_idx_list, flatten_shape))
  num_unique_emb_idxs = tf.shape(batch_unique_emb_idxs)[0]
  paddings = tf.zeros((max_num_batch_emb - num_unique_emb_idxs), dtype=tf.int32)
  # <int>[max_num_batch_emb_idxs]
  batch_emb_idxs = tf.concat((batch_unique_emb_idxs, paddings), axis=0)
  batch_emb_idxs = tf.where(batch_emb_idxs == -1, 0, batch_emb_idxs)
  batch_emb_idx_list = tf.where(
      tf.gather(batch_unique_emb_idxs, batch_emb_idx_list) == -1, -1,
      batch_emb_idx_list)
  # <int>[batch_size, num_nts, num_nodes]
  batch_emb_idx_list = tf.reshape(batch_emb_idx_list, tf.shape(emb_idx_list))
  return batch_emb_idxs, batch_emb_idx_list
Get unique embedding idxs within the batch.
language/compgen/csl/model/training/training_utils.py