Introduction

Batchwise indexing in TensorFlow means each element in the batch uses its own index list instead of sharing one global set of indices. This pattern appears in ranking models, beam search, candidate filtering, and sequence post-processing. The reliable solution is usually to combine tf.argsort or tf.math.top_k with tf.gather and set batch_dims correctly.

Understand the Shapes Before Indexing

Suppose you have a tensor with shape [batch, items, features]. If you gather along axis=1, you are selecting item rows for each batch entry.

1import tensorflow as tf
2
3payload = tf.constant([
4    [[101.0, 1.0], [102.0, 2.0], [103.0, 3.0]],
5    [[201.0, 4.0], [202.0, 5.0], [203.0, 6.0]],
6])
7
8scores = tf.constant([
9    [0.2, 0.9, 0.1],
10    [0.8, 0.4, 0.7],
11])
12
13print(payload.shape)  # (2, 3, 2)
14print(scores.shape)   # (2, 3)

The first dimension is batch size. The second dimension lines up between payload and scores. That alignment is what lets you sort scores and then reindex the payload with the same order.

Use tf.gather With batch_dims

TensorFlow's tf.gather supports per-batch indexing through batch_dims. In the TensorFlow API docs, batch_dims=1 is described as equivalent to looping over the first axis and gathering independently inside each batch row.

1indices = tf.constant([
2    [1, 0],
3    [2, 1],
4])
5
6selected = tf.gather(payload, indices, axis=1, batch_dims=1)
7print(selected)

This produces two selected item rows for each batch entry. Without batch_dims=1, TensorFlow interprets the indices differently and the result is usually not what you intended.

For batchwise work, think of batch_dims=1 as saying: "the first dimension of payload and indices already matches, so gather separately inside each batch row."

Sort Per Batch and Reorder a Paired Tensor

A very common workflow is: sort scores, get the order indices, then apply those indices to another tensor.

1order = tf.argsort(scores, axis=1, direction="DESCENDING")
2print(order)
3
4payload_sorted = tf.gather(payload, order, axis=1, batch_dims=1)
5print(payload_sorted)

This keeps the payload aligned with the score ordering. If you sort scores and forget to reorder the paired tensor with the same indices, the data becomes silently misaligned.

This is one of the most important habits in ranking code: treat the index tensor as the source of truth and reuse it for every related tensor that must stay synchronized.

Use tf.math.top_k When You Only Need the Best Results

If you need only the top k items instead of a full sort, tf.math.top_k is more direct.

1values, top_idx = tf.math.top_k(scores, k=2)
2print(values)
3print(top_idx)
4
5top_payload = tf.gather(payload, top_idx, axis=1, batch_dims=1)
6print(top_payload)

This is a common pattern for recommendation models and retrieval systems, where only the best few candidates matter.

When tf.gather_nd Is the Better Tool

tf.gather_nd is useful when you want explicit coordinate-based indexing rather than selecting along one axis.

1coords = tf.constant([
2    [0, 1],
3    [1, 2],
4])
5
6rows = tf.gather_nd(scores, coords)
7print(rows)

For straightforward batchwise row selection, tf.gather plus batch_dims is usually easier to read. Use gather_nd when the indexing logic genuinely depends on full coordinate tuples.

Add Shape Checks in Reusable Code

Batch indexing bugs often come from shape drift. Small assertions prevent long debugging sessions.

1tf.debugging.assert_rank(scores, 2)
2tf.debugging.assert_rank(payload, 3)
3tf.debugging.assert_equal(tf.shape(scores)[0], tf.shape(payload)[0])
4tf.debugging.assert_equal(tf.shape(scores)[1], tf.shape(payload)[1])

These checks are especially useful when the tensors come from different preprocessing steps.

Common Pitfalls

The most common mistake is forgetting batch_dims. That turns a per-batch operation into a global gather with different semantics.

Another issue is sorting one tensor and not applying the same index order to the related tensors. That produces wrong results without necessarily raising an error.

Developers also often choose the wrong axis after reshaping or batching logic changes. Confirm the tensor layout before writing the gather.

Finally, use tf.math.top_k when you only need a small best subset. A full sort works, but it does more work than necessary.

Summary

• Use tf.gather(..., axis=1, batch_dims=1) for per-batch indexing along the item dimension.
• Use tf.argsort to get a full order and tf.math.top_k for top results only.
• Reuse the same index tensor for every payload tensor that must stay aligned.
• Prefer tf.gather_nd only when you need explicit coordinate indexing.
• Add shape assertions so batchwise indexing errors fail early instead of silently.