Introduction

Boolean indexing in TensorFlow is the standard way to select or filter tensor values based on conditions. It is heavily used in preprocessing, masking losses, and extracting valid samples in sequence models. The most important tools are tf.boolean_mask, comparison operations, and tf.where for index-based selection.

Building Boolean Masks

A boolean mask is simply a tensor of True and False values. You usually create it by comparing a tensor to a threshold or category.

1import tensorflow as tf
2
3x = tf.constant([3, 8, 1, 10, 5], dtype=tf.int32)
4mask = x > 4
5
6print(mask.numpy())
7# [False  True False  True  True]

Masks can be chained using logical operators:

mask2 = tf.logical_and(x > 2, x < 9)
print(mask2.numpy())
# [ True  True False False  True]

This approach is vectorized and much faster than Python loops.

Filtering Values with tf.boolean_mask

Use tf.boolean_mask to keep values where mask entries are true.

1import tensorflow as tf
2
3x = tf.constant([3, 8, 1, 10, 5], dtype=tf.int32)
4mask = x > 4
5filtered = tf.boolean_mask(x, mask)
6
7print(filtered.numpy())
8# [ 8 10  5]

For two-dimensional tensors, mask the first dimension by default.

1m = tf.constant([
2    [1.0, 2.0],
3    [3.0, 4.0],
4    [5.0, 6.0]
5])
6row_mask = tf.constant([True, False, True])
7
8rows = tf.boolean_mask(m, row_mask)
9print(rows.numpy())
10# [[1. 2.]
11#  [5. 6.]]

Masking a Specific Axis

When needed, pass axis to apply mask on another dimension.

1import tensorflow as tf
2
3m = tf.constant([
4    [1, 2, 3],
5    [4, 5, 6]
6])
7col_mask = tf.constant([True, False, True])
8
9selected_cols = tf.boolean_mask(m, col_mask, axis=1)
10print(selected_cols.numpy())
11# [[1 3]
12#  [4 6]]

Axis-aware masking is useful for selecting valid channels or feature positions.

Using tf.where for Indices

Sometimes you need indices rather than filtered values. tf.where returns coordinates of true entries.

1import tensorflow as tf
2
3x = tf.constant([3, 8, 1, 10, 5], dtype=tf.int32)
4mask = x > 4
5idx = tf.where(mask)
6
7print(idx.numpy())
8# [[1]
9#  [3]
10#  [4]]
11
12picked = tf.gather(x, tf.squeeze(idx, axis=1))
13print(picked.numpy())
14# [ 8 10  5]

This pattern integrates well with custom indexing logic.

Practical ML Example: Masked Loss

Boolean masks are common when padding sequences and computing loss only on valid tokens.

1import tensorflow as tf
2
3# logits shape: batch, time, classes
4logits = tf.constant([
5    [[2.0, 0.5], [0.2, 1.8], [0.1, 0.0]],
6    [[1.2, 0.4], [0.0, 0.0], [0.0, 0.0]]
7])
8labels = tf.constant([
9    [0, 1, 0],
10    [0, 0, 0]
11], dtype=tf.int32)
12valid = tf.constant([
13    [True, True, True],
14    [True, False, False]
15])
16
17loss_per_token = tf.keras.losses.sparse_categorical_crossentropy(labels, logits, from_logits=True)
18masked_loss = tf.boolean_mask(loss_per_token, valid)
19mean_loss = tf.reduce_mean(masked_loss)
20
21print(mean_loss.numpy())

Without masking, padded positions can distort training signals.

Shape and Type Considerations

Key constraints to remember:

• Mask dtype must be boolean.
• Mask shape must align with selected axis.
• Result shape is usually dynamic because true count can vary.

For graph mode and model exporting, dynamic output length can affect downstream layers. If fixed shape is required, you may need padding after masking.

Performance Notes

Boolean masking is efficient, but repeated masking in tight loops can still create overhead. For heavy pipelines:

• Build masks once per batch when possible.
• Prefer fused tensor operations over Python control flow.
• Profile with realistic batch sizes.

In distributed training, deterministic mask logic is important to keep per-replica behavior consistent.

Common Pitfalls

• Using integer mask values instead of bool. Fix by producing mask from comparisons or casting to tf.bool.
• Mask shape mismatch on higher-rank tensors. Fix by validating rank and axis alignment.
• Expecting fixed output shape after masking. Fix by handling dynamic lengths explicitly.
• Confusing tf.where output with masked values. Fix by using tf.gather after index extraction.
• Forgetting to mask padded tokens in sequence losses. Fix by applying boolean mask before reduction.

Summary

• Boolean indexing in TensorFlow is primarily done with tf.boolean_mask.
• Build masks from vectorized comparisons and logical operations.
• Use axis for dimension-specific filtering and tf.where for indices.
• Masking is critical for correct loss computation in padded sequence tasks.
• Always validate mask dtype, shape, and downstream shape expectations.