Introduction

When people ask for NumPy-style 2D indexing in TensorFlow, they usually mean selecting elements using explicit row-column coordinate pairs. Simple slicing works for rectangular ranges, but coordinate-based element lookup is a different operation. In TensorFlow, the usual answer is tf.gather_nd, which is designed for exactly this style of indexed access.

Use tf.gather_nd for Coordinate Pairs

Suppose you have a matrix and want elements at positions (0, 1) and (2, 2).

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

This prints [11 32].

Each row in indices is one coordinate pair. That is the direct TensorFlow equivalent of NumPy-style advanced indexing by position list.

Understand the Shape of the Index Tensor

For a rank-2 tensor, each index row must have length 2, because you need one coordinate for the row and one for the column.

General rule:

• outer dimension counts how many selections you want
• inner dimension gives the coordinates into the tensor

For example, with three selections from a 2D tensor:

1indices = tf.constant([
2    [0, 0],
3    [1, 2],
4    [2, 1],
5])

This asks for three separate elements.

Compare It with NumPy

The NumPy equivalent often looks like this:

1import numpy as np
2
3matrix = np.array([
4    [10, 11, 12],
5    [20, 21, 22],
6    [30, 31, 32],
7])
8
9rows = np.array([0, 2])
10cols = np.array([1, 2])
11
12print(matrix[rows, cols])

In TensorFlow, you generally combine those row-column pairs into one 2D index tensor and use tf.gather_nd.

Use Slicing for Rectangular Regions

If the goal is not coordinate-based lookup but a rectangular block, ordinary slicing is simpler.

submatrix = matrix[0:2, 1:3]
print(submatrix.numpy())

This returns a slice, not an arbitrary list of individual elements. That distinction matters:

• slicing is for contiguous ranges
• 'gather_nd is for explicit coordinates'

Selecting Entire Rows or Columns Is Different Again

If you want complete rows or columns rather than point coordinates, tf.gather is often more direct.

1rows = tf.gather(matrix, [0, 2], axis=0)
2cols = tf.gather(matrix, [0, 2], axis=1)
3
4print(rows.numpy())
5print(cols.numpy())

This is different from NumPy-style pair indexing because it gathers along one axis at a time instead of selecting element-by-element coordinate pairs.

Batch and Higher-Rank Cases Still Use the Same Idea

tf.gather_nd also works for higher-rank tensors. You simply provide longer coordinate rows.

1tensor = tf.constant([
2    [[1, 2], [3, 4]],
3    [[5, 6], [7, 8]],
4])
5
6indices = tf.constant([
7    [0, 1, 0],
8    [1, 0, 1],
9])
10
11print(tf.gather_nd(tensor, indices).numpy())

Each index row now has length 3 because the tensor rank is 3.

Common Pitfalls

• Using normal slicing when the task actually needs arbitrary coordinate-based selection.
• Passing row and column lists separately instead of building one 2D index tensor for tf.gather_nd.
• Forgetting that each coordinate row must match the tensor rank being indexed.
• Reaching for tf.gather when the requirement is element-by-element coordinate lookup rather than axis-wise selection.
• Misreading the result shape because gather_nd returns one output entry per coordinate row.

Summary

• For NumPy-style coordinate indexing in TensorFlow, tf.gather_nd is usually the correct tool.
• Each row of the index tensor represents one coordinate into the tensor.
• Use ordinary slicing for rectangular ranges and tf.gather for row- or column-wise selection.
• The same pattern extends naturally to higher-rank tensors.
• The main job is to match the shape of the index tensor to the kind of indexing you actually want.