tf.ragged.boolean_mask

TensorFlow 2 version View source on GitHub

Applies a boolean mask to data without flattening the mask dimensions.

tf.ragged.boolean_mask(
    data, mask, name=None
)

Returns a potentially ragged tensor that is formed by retaining the elements in data where the corresponding value in mask is True.

  • output[a1...aA, i, b1...bB] = data[a1...aA, j, b1...bB]

    Where j is the ith True entry of mask[a1...aA].

Note that output preserves the mask dimensions a1...aA; this differs from tf.boolean_mask, which flattens those dimensions.

data A potentially ragged tensor.
mask A potentially ragged boolean tensor. mask's shape must be a prefix of data's shape. rank(mask) must be known statically.
name A name prefix for the returned tensor (optional).

A potentially ragged tensor that is formed by retaining the elements in data where the corresponding value in mask is True.

  • rank(output) = rank(data).
  • output.ragged_rank = max(data.ragged_rank, rank(mask) - 1).

ValueError if rank(mask) is not known statically; or if mask.shape is not a prefix of data.shape.

Examples:

  >>> # Aliases for True & False so data and mask line up.
  >>> T, F = (True, False)

<pre class="devsite-click-to-copy prettyprint lang-py">
  <code class="devsite-terminal" data-terminal-prefix="&gt;&gt;&gt;">tf.ragged.boolean_mask(  # Mask a 2D Tensor.</code>
  <code class="devsite-terminal" data-terminal-prefix="...">    data=[[1, 2, 3], [4, 5, 6], [7, 8, 9]],</code>
  <code class="devsite-terminal" data-terminal-prefix="...">    mask=[[T, F, T], [F, F, F], [T, F, F]]).tolist()</code>
  <code class="no-select nocode">  [[1, 3], [], [7]]</code>
  <code class="no-select nocode">  </code>
</pre>


<pre class="devsite-click-to-copy prettyprint lang-py">
  <code class="devsite-terminal" data-terminal-prefix="&gt;&gt;&gt;">tf.ragged.boolean_mask(  # Mask a 2D RaggedTensor.</code>
  <code class="devsite-terminal" data-terminal-prefix="...">    tf.ragged.constant([[1, 2, 3], [4], [5, 6]]),</code>
  <code class="devsite-terminal" data-terminal-prefix="...">    tf.ragged.constant([[F, F, T], [F], [T, T]])).tolist()</code>
  <code class="no-select nocode">  [[3], [], [5, 6]]</code>
  <code class="no-select nocode">  </code>
</pre>


<pre class="devsite-click-to-copy prettyprint lang-py">
  <code class="devsite-terminal" data-terminal-prefix="&gt;&gt;&gt;">tf.ragged.boolean_mask(  # Mask rows of a 2D RaggedTensor.</code>
  <code class="devsite-terminal" data-terminal-prefix="...">    tf.ragged.constant([[1, 2, 3], [4], [5, 6]]),</code>
  <code class="devsite-terminal" data-terminal-prefix="...">    tf.ragged.constant([True, False, True])).tolist()</code>
  <code class="no-select nocode">  [[1, 2, 3], [5, 6]]</code>
  <code class="no-select nocode">  </code>
</pre>