Introduction

Using one tensor to index another is a normal TensorFlow operation, but the exact API depends on what kind of indexing you need. For simple row or element selection, use tf.gather; for multi-dimensional coordinates, use tf.gather_nd; and for condition-based filtering, use tf.boolean_mask.

Indexing with tf.gather

tf.gather is the closest TensorFlow equivalent to selecting elements by integer index.

1import tensorflow as tf
2
3values = tf.constant([10, 20, 30, 40, 50])
4indices = tf.constant([0, 2, 4])
5
6selected = tf.gather(values, indices)
7print(selected.numpy())

Output:

[10 30 50]

This is the right choice when you have one tensor of positions and want items from another tensor along a single axis.

Gathering Rows from a Matrix

The same function works for higher-rank tensors. By default, it gathers along axis 0.

1matrix = tf.constant([
2    [1, 2],
3    [3, 4],
4    [5, 6],
5    [7, 8],
6])
7
8row_ids = tf.constant([3, 1])
9rows = tf.gather(matrix, row_ids)
10print(rows.numpy())

Output:

[[7 8]
 [3 4]]

If you want to gather columns instead, change the axis:

column_ids = tf.constant([1, 0])
cols = tf.gather(matrix, column_ids, axis=1)
print(cols.numpy())

Multi-Dimensional Coordinates with tf.gather_nd

If each index is a full coordinate, use tf.gather_nd.

1matrix = tf.constant([
2    [10, 11, 12],
3    [20, 21, 22],
4    [30, 31, 32],
5])
6
7coords = tf.constant([
8    [0, 1],
9    [2, 2],
10    [1, 0],
11])
12
13result = tf.gather_nd(matrix, coords)
14print(result.numpy())

Output:

[11 32 20]

Each row in coords points to one element in the source tensor.

Filtering with tf.boolean_mask

When the index tensor is boolean rather than integer, use tf.boolean_mask.

1values = tf.constant([5, 12, 7, 19, 3])
2mask = values > 8
3
4filtered = tf.boolean_mask(values, mask)
5print(filtered.numpy())

Output:

[12 19]

This is often the cleanest way to express "select all entries matching a condition."

Why Plain Python Indexing Is Not Always Enough

In eager mode, some simple indexing expressions work naturally:

values = tf.constant([10, 20, 30, 40])
print(values[2].numpy())

But for graph-friendly, batched, or dynamic index tensors, the TensorFlow gather APIs are more explicit and more reliable. They also map cleanly to GPU-accelerated kernels and shape inference.

Choosing the Right Operation

Use this rule of thumb:

• 'tf.gather for integer indices along one axis'
• 'tf.gather_nd for full multi-axis coordinates'
• 'tf.boolean_mask for boolean selection'

That distinction usually removes the confusion.

Batch Shapes Matter

Indexing bugs in TensorFlow are often shape bugs in disguise. If your source tensor is batched, confirm whether you want to index inside each batch element or across the batch axis itself. Printing tensor.shape before the gather step usually saves time because tf.gather and tf.gather_nd are both strict about how index shapes map to result shapes.

When the index values are produced by a model step such as tf.argmax, it is also worth checking the dtype. Gather operations expect integer index tensors, so accidental casting to floating point will fail even if the numeric values look correct.

Common Pitfalls

• Passing floating-point indices instead of integer tensors.
• Using tf.gather when the index tensor actually contains full coordinates and tf.gather_nd is required.
• Forgetting the axis argument and gathering along the wrong dimension.
• Expecting boolean masks and integer indices to behave like the same operation.

Summary

• TensorFlow supports tensor-based indexing directly, but through specific APIs.
• Use tf.gather for standard integer indexing.
• Use tf.gather_nd for coordinate-based selection.
• Use tf.boolean_mask for condition-based filtering.
• Picking the right indexing primitive is the main step; the rest is mostly shape management.