def boolean_mask(tensor, mask):
  """Select elements from tensor where mask is True, and pads the rest with 0.

  Example:
    tensor = [1 2 3 4 5 6 7]
    mask   = [T T F F T T T]
    output = [1 2 5 6 7 0 0]

  Args:
    tensor: <T> [batch_size, dim, ...]
    mask: <bool> [batch_size, dim]

  Returns:
    masked_tensor: <T> [batch_size, dim, ...]. The first k items of row i
      correspond to the k truthy values in row i of the mask tensor,
      indexed into the provided tensor.
  """
  batch_size = shape(mask, 0)
  dim = shape(mask, 1)

  # Bring the coordinates of un-masked elements to the front.
  mask = tf.cast(mask, tensor.dtype)
  indices = tf.argsort(mask, axis=1, direction="DESCENDING", stable=True)
  batch_indices = tf.expand_dims(tf.range(tf.shape(mask)[0]), 1)
  batch_indices = tf.tile(batch_indices, [1, tf.shape(indices)[1]])
  coords = tf.stack([batch_indices, indices], axis=2)

  # Gather and zero masked elements.
  selected_tensor = tf.gather_nd(tensor, coords)
  selected_mask = tf.gather_nd(mask, coords)
  reshape = [batch_size, dim] + [1] * (tensor.get_shape().ndims - 2)
  selected_mask = tf.reshape(selected_mask, reshape)
  masked_tensor = selected_tensor * selected_mask

  return masked_tensor
Select elements from tensor where mask is True, and pads the rest with 0.

  Example:
    tensor = [1 2 3 4 5 6 7]
    mask   = [T T F F T T T]
    output = [1 2 5 6 7 0 0]

  Args:
    tensor: <T> [batch_size, dim, ...]
    mask: <bool> [batch_size, dim]

  Returns:
    masked_tensor: <T> [batch_size, dim, ...]. The first k items of row i
      correspond to the k truthy values in row i of the mask tensor,
      indexed into the provided tensor.
  
language/capwap/utils/tensor_utils.py