def _convert_ids_to_onehot(input_ids, input_mask, vocab_size):
  """Returns a k-hot vector for the given input ids in vocab.

  Args:
    input_ids: (tf.int32) batch_size x seq_length
    input_mask: (tf.int32) batch_size x seq_length
    vocab_size: scalar

  Returns:
    input_onehot: (tf.int32) batch_size x vocab_size
  """
  onehot = tf.one_hot(input_ids, vocab_size, dtype=tf.int32, axis=-1)
  masked_onehot = onehot * tf.expand_dims(input_mask, axis=2)
  input_onehot = tf.minimum(1, tf.reduce_sum(masked_onehot, axis=1))
  return input_onehot
Returns a k-hot vector for the given input ids in vocab.

  Args:
    input_ids: (tf.int32) batch_size x seq_length
    input_mask: (tf.int32) batch_size x seq_length
    vocab_size: scalar

  Returns:
    input_onehot: (tf.int32) batch_size x vocab_size
  
language/labs/drkit/run_dualencoder_lsf.py