def ensure_values_in_mat(mat, values, dtype):
  """Ensure indices are present in the rows of mat.

  This replaces the last value in each row of mat with the corresponding
  value in values, IF that value is not in any of the remaining columns.

  Args:
    mat: [batch_size, M] Tensor.
    values: [batch_size] Tensor.
    dtype: Data type of mat.

  Returns:
    [batch_size, M] Tensor.
  """
  val_in_rows = tf.reduce_sum(
      tf.cast(tf.equal(mat[:, :-1], tf.expand_dims(values, 1)), dtype), axis=1)
  last_col = val_in_rows * mat[:, -1] + (1 - val_in_rows) * values
  return tf.concat((mat[:, :-1], tf.expand_dims(last_col, 1)), axis=1)
Ensure indices are present in the rows of mat.

  This replaces the last value in each row of mat with the corresponding
  value in values, IF that value is not in any of the remaining columns.

  Args:
    mat: [batch_size, M] Tensor.
    values: [batch_size] Tensor.
    dtype: Data type of mat.

  Returns:
    [batch_size, M] Tensor.
  
language/labs/drkit/model_fns.py