Tensor Indexing, Slicing and tf.gather in TensorFlow

The Story — Picking Seats in a Cinema

Getting parts of a tensor is something you will do every day. You will take one image from a batch, the last time step of a sequence, or the rows of a word-embedding table. This lesson covers all the main tools.

Indexing a 1-D Tensor

Indexing works just like Python lists. The first position is 0. Negative numbers count from the end.

import tensorflow as tf

x = tf.constant([10, 20, 30, 40, 50, 60])

print(x[0])        # first
print(x[-1])       # last
print(x[2].numpy(), x[-2].numpy())

Slicing — start : stop : step

A slice has three parts. Each part is optional.

x = tf.constant([10, 20, 30, 40, 50, 60])

print(x[1:4].numpy())     # positions 1, 2, 3
print(x[:3].numpy())      # first three
print(x[3:].numpy())      # from 3 to the end
print(x[::2].numpy())     # every second value
print(x[::-1].numpy())    # reversed
print(x[-3:].numpy())     # last three

Indexing Many Dimensions

For a tensor with more axes, separate the index for each axis with a comma: x[row, col].

m = tf.reshape(tf.range(12), [3, 4])
print(m)

print("m[1, 2]  =", m[1, 2].numpy())            # one value
print("m[1]     =", m[1].numpy())               # whole row 1
print("m[:, 2]  =", m[:, 2].numpy())            # whole column 2
print("m[0:2, 1:3] =\n", m[0:2, 1:3].numpy())   # a block

The Ellipsis (...) and tf.newaxis

... means "all the axes in between". tf.newaxis adds an axis of size 1.

images = tf.zeros([32, 28, 28, 3])       # batch, height, width, channels

print(images[..., 0].shape)            # red channel of every image
print(images[0, ...].shape)            # first image
print(images[:, :, :, tf.newaxis].shape)

Try It — The Slice Playground

You Cannot Assign Into a Tensor

Tensors are immutable. The line x[0] = 5 fails. You have two options.

x = tf.constant([1, 2, 3, 4])

# Option 1: build a NEW tensor with some values replaced
y = tf.tensor_scatter_nd_update(x, indices=[[0], [2]], updates=[100, 300])
print(y.numpy())

# Option 2: use a Variable, which can change in place
v = tf.Variable([1, 2, 3, 4])
v[1:3].assign([20, 30])
print(v.numpy())

tf.gather — Pick Any Positions

Slices can only pick evenly spaced positions. tf.gather(params, indices, axis=0) picks any list of positions, in any order. You can even repeat a position.

m = tf.constant([[1, 2, 3],
                 [4, 5, 6],
                 [7, 8, 9],
                 [10, 11, 12]])

print(tf.gather(m, [3, 0, 0]))            # rows 3, 0, 0
print(tf.gather(m, [2, 0], axis=1))       # columns 2 and 0

vocab_size, dim = 6, 3
table = tf.reshape(tf.range(vocab_size * dim, dtype=tf.float32), [vocab_size, dim])

sentence = tf.constant([4, 1, 4, 2])        # "the cat the dog" as word IDs
vectors = tf.gather(table, sentence)
print(vectors.shape)
print(vectors.numpy())

tf.gather_nd — Pick Single Cells by Full Address

With gather_nd, each index is a full address like [row, col]. You get one value per address.

cells = tf.gather_nd(m, indices=[[0, 0], [1, 2], [3, 1]])
print(cells.numpy())       # m[0,0], m[1,2], m[3,1]

Try It — The Gather Lab

Picking by Condition — boolean_mask and tf.where

scores = tf.constant([[55, 91, 38],
                      [72, 45, 88]])

passed = scores >= 50                     # a bool tensor, same shape
print(passed.numpy())

print(tf.boolean_mask(scores, passed).numpy())   # values that passed (flattened)
print(tf.where(passed).numpy())                  # [row, col] of each True
print(tf.where(passed, scores, 0).numpy())       # keep passed, else 0

Which Tool Should I Use?

Golden Rules