def batch_boolean_mask(mask):
  """Get indices of true values.

  Args:
    mask: [batch_size, num_values]

  Returns:
    true_indices: [batch_size, max_true]
    gathered_mask: [batch_size, max_true]
  """
  # [batch_size, num_values]
  mask = tf.to_int32(mask)

  # [batch_size]
  num_true = tf.reduce_sum(mask, 1)

  # []
  max_true = tf.reduce_max(num_true)

  # [batch_size, max_true]
  gathered_mask, true_indices = tf.nn.top_k(mask, max_true)
  gathered_mask = tf.cast(gathered_mask, tf.bool)

  return gathered_mask, true_indices
Get indices of true values.

  Args:
    mask: [batch_size, num_values]

  Returns:
    true_indices: [batch_size, max_true]
    gathered_mask: [batch_size, max_true]
  
language/labs/exemplar_decoding/utils/tensor_utils.py