Introduction

tf.gather_nd selects values or slices from a tensor using multi-dimensional index coordinates. The intuitive way to think about it is "advanced indexing with coordinate tuples". It is especially useful when plain tf.gather is not enough because you need to index across multiple axes at once.

Core Sections

Mental Model for gather_nd

With tf.gather, you select entries along one axis. With tf.gather_nd, each index row points to one coordinate in the source tensor, or to a slice if the index depth is smaller than tensor rank.

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

Each index pair selects one scalar value from a two-dimensional tensor.

Gather Whole Slices with Partial Coordinates

If index depth is less than rank, gather_nd returns sub-tensors.

1import tensorflow as tf
2
3params = tf.constant([
4    [[1, 2], [3, 4]],
5    [[5, 6], [7, 8]],
6], dtype=tf.int32)  # shape [2, 2, 2]
7
8indices = tf.constant([
9    [0],
10    [1],
11])
12
13out = tf.gather_nd(params, indices)
14print(out.numpy())
15# [[[1 2]
16#   [3 4]]
17#  [[5 6]
18#   [7 8]]]

Here each coordinate chooses one first-axis block.

Compare with NumPy Advanced Indexing

If you know NumPy, gather_nd is close to indexing by explicit coordinate matrix. It is often used in sequence models, object detection pipelines, and custom losses where selected positions differ per example.

1import tensorflow as tf
2
3batch_logits = tf.constant([
4    [0.1, 0.8, 0.1],
5    [0.7, 0.2, 0.1],
6    [0.2, 0.2, 0.6],
7], dtype=tf.float32)
8labels = tf.constant([1, 0, 2], dtype=tf.int32)
9
10row_ids = tf.range(tf.shape(labels)[0], dtype=tf.int32)
11idx = tf.stack([row_ids, labels], axis=1)
12true_class_scores = tf.gather_nd(batch_logits, idx)
13print(true_class_scores.numpy())

This extracts one score per row according to label.

Shape Rule to Remember

Output shape is indices-shape-prefix plus params-shape-suffix after indexed dimensions. If this sounds abstract, inspect shapes with small tensors during development.

print("params shape:", params.shape)
print("indices shape:", indices.shape)
print("output shape:", out.shape)

Consistent shape debugging prevents most gather_nd errors.

Performance and Readability Guidance

gather_nd is powerful but can make code harder to read. Build indices step by step with clear variable names, and add shape assertions around custom logic. For repeated patterns, wrap index creation in helper functions.

Real Workflow Example with Batched Coordinates

A common pattern is selecting values from batched tensors where each batch row has a different target coordinate. gather_nd handles this cleanly when you build row indices and target indices together.

1import tensorflow as tf
2
3# shape: [batch, seq_len]
4scores = tf.constant([
5    [0.1, 0.7, 0.2],
6    [0.8, 0.1, 0.1],
7    [0.2, 0.3, 0.5],
8], dtype=tf.float32)
9
10best_pos = tf.constant([1, 0, 2], dtype=tf.int32)
11rows = tf.range(tf.shape(scores)[0], dtype=tf.int32)
12coords = tf.stack([rows, best_pos], axis=1)
13chosen = tf.gather_nd(scores, coords)
14print(chosen.numpy())

This avoids loops and keeps graph operations vectorized. It is widely used in sequence labeling and custom decoding logic.

When debugging index logic, test with tiny tensors and print coordinate arrays directly. Clear intermediate outputs reduce shape-related mistakes quickly.

Readable helper functions for index construction make complex model code easier to review and maintain.

Consistent tensor-shape documentation also speeds onboarding for new contributors.

Use small reproducible examples when reviewing indexing bugs.

Common Pitfalls

• Using tf.gather_nd where simple tf.gather would be clearer.
• Building indices with wrong dtype instead of integer tensors.
• Misunderstanding index depth and getting unexpected slice outputs.
• Ignoring shape rules and debugging only after runtime failures.
• Creating complicated one-liners that hide index construction logic.

Summary

• tf.gather_nd selects tensor values using coordinate tuples.
• It can return scalars or slices depending on index depth.
• It is ideal for per-row or per-example position selection tasks.
• Validate shapes and index dtypes early.
• Prefer readable index-building code for maintainable models.