Introduction

Selecting values from each row of a tensor comes up constantly in TensorFlow code: feature picking, top-k post-processing, masked retrieval, and custom losses all use it. The correct API depends on the shape of your indices. If every row needs the same columns, plain tf.gather is enough. If each row has its own column choices, you usually want either tf.gather_nd or tf.gather with batch_dims=1.

Same columns for every row: use tf.gather

When all rows need the same column positions, the job is simple. Gather along axis=1.

1import tensorflow as tf
2
3x = tf.constant([
4    [10, 11, 12, 13],
5    [20, 21, 22, 23],
6    [30, 31, 32, 33],
7], dtype=tf.int32)
8
9cols = tf.constant([1, 3], dtype=tf.int32)
10selected = tf.gather(x, cols, axis=1)
11
12print(selected.numpy())

This returns a tensor of shape rows x len(cols). It is the cleanest option because the index list is shared across the batch.

One different column per row: use tf.gather_nd

If each row points to a different column, build explicit row-column pairs and pass them to tf.gather_nd.

1import tensorflow as tf
2
3x = tf.constant([
4    [5, 6, 7],
5    [8, 9, 10],
6    [11, 12, 13],
7], dtype=tf.int32)
8
9col_index = tf.constant([2, 0, 1], dtype=tf.int32)
10row_index = tf.range(tf.shape(x)[0], dtype=tf.int32)
11indices = tf.stack([row_index, col_index], axis=1)
12
13picked = tf.gather_nd(x, indices)
14print(picked.numpy())

The result is one value per row: from row 0 take column 2, from row 1 take column 0, and from row 2 take column 1.

This pattern is easy to reason about because you make the coordinates explicit.

Multiple different columns per row: use batch_dims=1

TensorFlow's batch_dims argument is often the nicest answer when every row has its own list of columns. It tells tf.gather to treat the leading dimension as the batch and gather inside each row separately.

1import tensorflow as tf
2
3x = tf.constant([
4    [10, 11, 12, 13],
5    [20, 21, 22, 23],
6    [30, 31, 32, 33],
7], dtype=tf.int32)
8
9per_row_cols = tf.constant([
10    [0, 2],
11    [1, 3],
12    [3, 0],
13], dtype=tf.int32)
14
15selected = tf.gather(x, per_row_cols, axis=1, batch_dims=1)
16print(selected.numpy())

The output is:

• row 0: columns 0 and 2
• row 1: columns 1 and 3
• row 2: columns 3 and 0

This is often easier to read than hand-building a rank-3 gather_nd index tensor.

Pay attention to shapes and dtypes

Most bugs here are shape bugs. During prototyping, print the tensor shapes and assert the index dtype before wiring the operation into a model.

1tf.debugging.assert_rank(x, 2)
2tf.debugging.assert_type(per_row_cols, tf.int32)
3
4print("x shape:", x.shape)
5print("index shape:", per_row_cols.shape)

TensorFlow expects integer indices, usually int32 or int64. It also helps to write a tiny test tensor by hand and compute the expected answer yourself before applying the same logic to a large model tensor.

Keep indexing vectorized

If you find yourself looping over rows in Python and slicing each row individually, that is usually a sign the code should be rewritten with tf.gather, tf.gather_nd, or batch_dims. Vectorized indexing keeps the computation inside TensorFlow, which is important for both performance and graph compilation.

This matters even more in tf.data pipelines or model code that may later run under tf.function, because Python loops can become a performance bottleneck or complicate tracing.

Common Pitfalls

The most common mistake is using tf.gather for row-specific indices without setting batch_dims. That gathers the same index pattern for every row and silently gives the wrong shape.

Another issue is building gather_nd indices with the wrong last dimension. For a rank-2 tensor, each index vector must contain exactly two integers: row and column.

Developers also mix up axis=0 and axis=1. If you are selecting columns from each row, axis=1 is the important axis.

Finally, do not rely on accidental behavior from invalid indices. Validate shapes and bounds early so indexing bugs fail near the source.

Summary

• Use tf.gather(..., axis=1) when every row needs the same columns.
• Use tf.gather_nd when you want explicit row-column coordinate pairs.
• Use tf.gather(..., batch_dims=1) when each row has its own list of column indices.
• Validate index dtype and shape before using the selection inside a model.
• Keep row-wise selection vectorized instead of looping in Python.