Indexing and slicing tensors in TensorFlow

Introduction

Indexing and slicing let you pick parts of your data easily, like choosing pieces of a photo or words from a sentence.

When you want to get a single value from a big data set, like one pixel from an image.

When you need to extract a smaller part of your data, like a few rows from a table.

When you want to change or analyze only a part of your data without touching the rest.

When preparing batches of data for training a machine learning model.

When you want to combine or compare specific parts of different data tensors.

Syntax

TensorFlow

tensor[index]
tensor[start:stop]
tensor[start:stop:step]
tensor[dim1_index, dim2_index, ...]

Indexing starts at 0, so the first element is at position 0.

You can use negative numbers to count from the end, like -1 for the last element.

Examples

Gets the third element (index 2) from a 1D tensor.

TensorFlow

import tensorflow as tf
x = tf.constant([10, 20, 30, 40, 50])
print(x[2])

Slices elements from index 1 up to but not including 4.

TensorFlow

x = tf.constant([10, 20, 30, 40, 50])
print(x[1:4])

Gets the element at row 1, column 0 from a 2D tensor.

TensorFlow

x = tf.constant([[1, 2], [3, 4], [5, 6]])
print(x[1, 0])

Reverses the tensor using slicing with a negative step.

TensorFlow

x = tf.constant([10, 20, 30, 40, 50])
print(x[::-1])

Sample Model

This program shows how to get a single element, slice a part of a 2D tensor, and reverse the rows of the tensor.

TensorFlow

import tensorflow as tf

# Create a 2D tensor (3 rows, 4 columns)
tensor = tf.constant([[1, 2, 3, 4],
                      [5, 6, 7, 8],
                      [9, 10, 11, 12]])

# Get the element at row 0, column 2
elem = tensor[0, 2]

# Slice rows 1 to 2 (inclusive start, exclusive end 3), columns 1 to 3
slice_part = tensor[1:3, 1:4]

# Reverse the rows
reversed_rows = tensor[::-1, :]

print(f"Element at [0,2]: {elem.numpy()}")
print("Slice rows 1-2, cols 1-3:")
print(slice_part.numpy())
print("Tensor with rows reversed:")
print(reversed_rows.numpy())

Important Notes

TensorFlow tensors are immutable, so slicing returns a new tensor without changing the original.

You can combine indexing and slicing to access complex parts of your data easily.

Using .numpy() converts a tensor to a regular array for printing or further use outside TensorFlow.

Summary

Indexing picks single elements by position.

Slicing extracts ranges of elements using start, stop, and step.

You can use indexing and slicing together on multi-dimensional tensors.