TensorFlow
indexing
tensor operations
machine learning
Python

How to index a list with a TensorFlow tensor?

Master System Design with Codemia

Enhance your system design skills with over 120 practice problems, detailed solutions, and hands-on exercises.

Introduction

A Python list and a TensorFlow tensor live in different execution models. If your index is a TensorFlow value, the safest approach is usually to convert the data to a tensor and use TensorFlow indexing operations such as tf.gather.

Why Direct Python List Indexing Is Fragile

A normal Python list expects a normal Python integer.

python
values = [10, 20, 30]

A TensorFlow tensor is not just an integer in another wrapper. It is a TensorFlow object that may be traced, staged into a graph, or computed later.

That means code like this is unreliable:

python
1import tensorflow as tf
2
3values = [10, 20, 30]
4index = tf.constant(1)
5# result = values[index]

In eager mode, you may be tempted to convert with .numpy() and move on. That works only when you are already on the Python host side and do not care about graph execution, tracing, or model portability.

The Right Fix: Convert The Data And Use tf.gather

If the index is a tensor, keep the data in TensorFlow as well.

python
1import tensorflow as tf
2
3values = tf.constant([10, 20, 30, 40, 50], dtype=tf.int32)
4indices = tf.constant([0, 2, 4], dtype=tf.int32)
5
6result = tf.gather(values, indices)
7print(result.numpy())

This works in eager mode and inside @tf.function, which is the real reason it is the preferred solution.

Single Index And Multiple Indices

tf.gather works for both one index and many indices.

python
1import tensorflow as tf
2
3values = tf.constant([100, 200, 300], dtype=tf.int32)
4
5single = tf.gather(values, tf.constant(1))
6many = tf.gather(values, tf.constant([2, 0]))
7
8print(single.numpy())
9print(many.numpy())

That makes it a good general replacement for host-side list indexing.

Multi-Dimensional Cases

When the source is a matrix or higher-dimensional tensor, you can gather along a chosen axis.

python
1import tensorflow as tf
2
3matrix = tf.constant([
4    [1, 2, 3],
5    [4, 5, 6],
6    [7, 8, 9],
7], dtype=tf.int32)
8
9rows = tf.gather(matrix, tf.constant([2, 0]), axis=0)
10cols = tf.gather(matrix, tf.constant([1, 2]), axis=1)
11
12print(rows.numpy())
13print(cols.numpy())

If you need coordinate-style lookup such as "row 0, col 1" and "row 2, col 2", use tf.gather_nd instead.

python
coords = tf.constant([[0, 1], [2, 2]], dtype=tf.int32)
points = tf.gather_nd(matrix, coords)
print(points.numpy())

What About .numpy()?

You can convert a scalar tensor to a Python integer in eager mode:

python
1index = tf.constant(1)
2python_index = int(index.numpy())
3values = [10, 20, 30]
4print(values[python_index])

This is acceptable for small host-side scripts. It is the wrong pattern inside model code, inside @tf.function, or anywhere you want TensorFlow to own the whole computation.

A useful rule is:

  • if the code is ordinary Python control flow, conversion may be okay
  • if the code belongs to a TensorFlow pipeline, stay in TensorFlow ops

Boolean And Conditional Selection

If the goal is filtering rather than positional lookup, tf.boolean_mask may be clearer than gather.

python
1import tensorflow as tf
2
3x = tf.constant([3, 10, 4, 20, 7], dtype=tf.int32)
4mask = x > 6
5filtered = tf.boolean_mask(x, mask)
6print(filtered.numpy())

This is often the better answer when the real problem is "pick elements meeting a condition."

Bounds And Dtypes Matter

TensorFlow indices should be integer tensors such as int32 or int64. Out-of-range indices will raise runtime errors.

python
1import tensorflow as tf
2
3values = tf.constant([10, 20, 30], dtype=tf.int32)
4indices = tf.constant([0, 5], dtype=tf.int32)
5# tf.gather(values, indices)  # out of range

If indices come from predictions or external input, validate them before gathering.

Common Pitfalls

The most common mistake is indexing a Python list with a tensor inside code that later moves into @tf.function. It may appear to work during experimentation and then fail when traced.

Another mistake is converting tensors to Python too early. That breaks TensorFlow's ability to optimize and stage the computation.

Developers also use tf.gather when they really need tf.gather_nd or tf.boolean_mask. Those operations solve different indexing problems.

Finally, check index dtype and range. Many indexing bugs are just invalid values, not TensorFlow syntax problems.

Summary

  • A TensorFlow tensor is not a drop-in Python list index in graph-oriented code.
  • Convert the source data to a tensor and use tf.gather for normal indexed lookup.
  • Use tf.gather_nd for coordinate lookup and tf.boolean_mask for filtering.
  • Converting with .numpy() is only appropriate for simple eager-mode host logic.
  • Keep indexing inside TensorFlow ops when building model or pipeline code.

Course illustration
Course illustration

All Rights Reserved.