Indexing, gathering and reshaping

TensorFlow · Essentials · 10 min

Preparing inputs, picking predictions and building masks are all indexing and reshaping tasks.

Basic slicing follows Python and NumPy: x[1], x[:, 1], x[1:, ::2], x[..., 0], and tf.newaxis (alias for None) to insert a dimension.

What NumPy does with fancy indexing, TF does with functions:

  • tf.gather(x, indices, axis=0) picks slices by index, e.g. rows 2 and 0.
  • tf.gather_nd(x, [[0, 1], [2, 3]]) picks individual elements by full coordinates.
  • tf.boolean_mask(x, mask) keeps elements where a boolean mask is true (returns a flat result when the mask covers all dims).
  • tf.where(cond, a, b) chooses element-wise between two tensors, with broadcasting.

Because tensors are immutable, "updating" a slice means building a new tensor, for example with tf.where or tf.tensor_scatter_nd_update.

Reshaping:

  • tf.reshape(x, [3, -1]) regroups elements (one -1 is inferred).
  • tf.expand_dims(x, 0) and tf.squeeze(x) add or remove size-1 dimensions.
  • tf.transpose(x, perm=[0, 2, 1]) reorders dimensions.
  • tf.concat(list, axis) joins along an existing axis; tf.stack(list, axis) adds a new one.

TensorFlow's image layers use channels-last layout (batch, height, width, channels) by default, unlike PyTorch.

Example

import tensorflow as tf

x = tf.reshape(tf.range(12), [3, 4])
print(x[:, 1].numpy())                        # [1 5 9]
print(x[1:, ::2].numpy())                     # [[ 4  6] [ 8 10]]
print(tf.gather(x, [2, 0])[:, 0].numpy())     # [8 0]   rows 2 and 0
print(tf.boolean_mask(x, x > 8).numpy())      # [ 9 10 11]
print(tf.where(x > 8, x, 0)[2].numpy())       # [ 0  9 10 11]

print(tf.expand_dims(x, 0).shape)             # (1, 3, 4)
print(tf.transpose(x).shape)                  # (4, 3)
print(tf.concat([x, x], axis=0).shape)        # (6, 4)
print(tf.stack([x, x]).shape)                 # (2, 3, 4)
print(tf.squeeze(tf.zeros([1, 3, 1])).shape)  # (3,)

Remember

  • Slicing like NumPy; tf.gather/boolean_mask/where for fancy selection
  • Updates create new tensors (tf.where, tensor_scatter_nd_update)
  • Images are channels-last: (N, H, W, C)

Watch out for

  • Expecting x[[2, 0]] to gather rows as in NumPy; use tf.gather.

Check yourself

  1. Which call picks rows 3 and 1 (in that order) of a 2-d tensor x?

Sign in to answer, see the explanations, and get this lesson back as a quick review later.

Official docs: www.tensorflow.org/guide/tensor_slicing