Introduction

TensorFlow supports a lot of indexing patterns that feel similar to NumPy, but the translation is not one-to-one. Simple slicing usually works as expected, while more advanced selection patterns are clearer and safer when expressed with TensorFlow operators such as tf.gather, tf.gather_nd, and tf.boolean_mask.

Use Basic Slicing First

Ordinary slice syntax is the most readable option when you only need straightforward row, column, or range access.

1import tensorflow as tf
2
3x = tf.constant([
4    [1, 2, 3],
5    [4, 5, 6],
6    [7, 8, 9],
7])
8
9print(x[0, 1].numpy())
10print(x[:, 1:].numpy())
11print(x[1].numpy())

If your code can be expressed this way, keep it this way. It is close to NumPy, easy to review, and works well in eager execution.

Use tf.gather for Indexed Selection Along One Axis

When the indices come from another tensor instead of being hard-coded in the source code, tf.gather is usually the right next step.

1import tensorflow as tf
2
3values = tf.constant([10, 20, 30, 40])
4indices = tf.constant([3, 1, 1, 0])
5print(tf.gather(values, indices).numpy())
6
7matrix = tf.constant([
8    [1, 2],
9    [3, 4],
10    [5, 6],
11])
12rows = tf.constant([2, 0])
13print(tf.gather(matrix, rows, axis=0).numpy())

The axis argument is important. A valid-looking gather on the wrong axis can silently produce the wrong shape and the wrong semantics.

Use tf.gather_nd and tf.boolean_mask for Advanced Cases

If you need coordinate-based lookup across multiple dimensions, tf.gather_nd is a better match than trying to force a NumPy-style expression into plain bracket syntax.

1import tensorflow as tf
2
3t = tf.constant([
4    [11, 12, 13],
5    [21, 22, 23],
6    [31, 32, 33],
7])
8coords = tf.constant([
9    [0, 2],
10    [2, 1],
11    [1, 0],
12])
13
14print(tf.gather_nd(t, coords).numpy())

For boolean filtering, use tf.boolean_mask.

1import tensorflow as tf
2
3x = tf.constant([5, 9, 2, 7, 1])
4mask = x > 4
5print(tf.boolean_mask(x, mask).numpy())

A useful mental model is to classify the problem before coding it: slice, gather on one axis, gather by coordinates, or mask by condition. Once you know which category you have, the right TensorFlow API is usually obvious.

Tensor Updates Need Scatter Ops

A common NumPy habit is in-place assignment. TensorFlow tensors are immutable, so indexed updates need scatter-style operations instead.

1import tensorflow as tf
2
3base = tf.constant([0, 0, 0, 0], dtype=tf.int32)
4indices = tf.constant([[1], [3]])
5updates = tf.constant([9, 5], dtype=tf.int32)
6
7result = tf.tensor_scatter_nd_update(base, indices, updates)
8print(result.numpy())

This matters a lot when porting older NumPy-heavy code into TensorFlow training or serving pipelines.

Keep the Work Inside TensorFlow

It is tempting to call .numpy(), perform indexing in NumPy, and convert back to a tensor. That works for quick experiments, but it is usually the wrong design for traced functions, accelerators, and performance-sensitive code. Staying inside TensorFlow keeps execution more portable and easier to optimize.

After complex indexing, shape checks are worth adding:

y = tf.gather(tf.random.normal((8, 16, 4)), [0, 3, 5], axis=0)
tf.debugging.assert_shapes([(y, (3, 16, 4))])

Shape assertions catch many indexing bugs earlier than silent downstream failures.

Common Pitfalls

Assuming every advanced NumPy indexing expression works unchanged in TensorFlow leads to brittle code. Translate the intent, not the exact syntax.

Using the wrong axis in tf.gather is an easy way to get valid but incorrect results.

Trying to update tensors in place fails because tensors are immutable. Use scatter update operations instead.

Summary

• Use ordinary slicing when the selection is simple and static.
• Use tf.gather for index tensors along one axis.
• Use tf.gather_nd for coordinate-style selection and tf.boolean_mask for conditional filtering.
• Use scatter update APIs when you need indexed writes.