Tensor Indexing in Tensorflow

Last Updated : 28 Apr, 2025

In the realm of machine learning and deep learning, tensors are fundamental data structures used to represent numerical data with multiple dimensions. TensorFlow, a powerful numerical computation library, equips you with an intuitive and versatile set of operations for manipulating and accessing data within these tensors. Understanding tensor indexing in TensorFlow becomes crucial for navigating data effectively and building machine learning models.

Tensor in TensorFlow

In TensorFlow, a tensor is a fundamental data structure used to represent multi-dimensional arrays or tensors. Tensors are the primary building blocks of TensorFlow computations, serving as the basic units for storing and manipulating data.

Tensors can be

Scalars (0-dimensional tensors)
Vectors (1-dimensional tensors)
Matrices (2-dimensional tensors)
higher-dimensional arrays.

Characteristics of tensors in TensorFlow

Shape: Tensors have a shape that determines their dimensions and size along each axis.
For example, A 2x3 matrix has the shape (2, 3), which represents two rows and three columns.
Data Type: Tensors have a data type, which describes the type of data stored in its elements. Common data types are float32, int32, and string.
Rank: The rank of a tensor indicates the number of dimensions it holds.
For example:
- A scalar tensor (a single value) has a rank of 0.
- A vector tensor (a 1-dimensional array) has a rank of 1.
- A matrix tensor (a 2-dimensional array) has a rank of 2
Axis: Axis refers to a dimension of a tensor. In TensorFlow, tensors can have multiple axes, each representing a different dimension of the data.
For Example:
- A 2D tensor has two axes: rows and columns (0 and 1).
- A 3D tensor has three axes (0 for height, 1 for width, and 2 for depth) (e.g., the first axis typically refers to rows in matrices).
Values: Tensors contain values or elements stored in a contiguous memory block. These values can be accessed and manipulated using various TensorFlow operations

Tensor Indexing in TensorFlow

Tensor indexing is the process of accessing and manipulating certain elements or subsets of a tensor. Tensor indexing, similar to the array indexing in Python, allows us to extract specific pieces or slices of data from a tensor. This is required for several TensorFlow operations, including data preparation, feature extraction, and model assessment.

Accessing or modifying specific elements within these tensors becomes crucial for effective manipulation. This process, known as tensor indexing, mirrors how you interact with lists or strings in Python or arrays in NumPy. TensorFlow offers specialized operations like tf.slice, tf.gather, tf.gather_nd, and tf.scatter_nd to empower you with various tensor indexing tasks. Delving deeper, we'll explore how to leverage these operations for slicing, extracting, and inserting data within your tensors.

This article explores various ways to index tensors:

1. Slicing

Slicing in TensorFlow lets you grab specific sections of your data, just like picking out a slice of cake! It's useful for extracting smaller vectors, matrices, or even higher-dimensional chunks from your tensors.

Slice It Up with tf.slice:

This operation takes three ingredients:

The main ingredient: the tensor you want to slice.
Starting point: where you want to begin the slice (like choosing your first bite).
Size: how much you want to take (a small, dainty slice or a generous piece?).

Imagine a delicious tensor t1 full of numbers:

Python3

import tensorflow as tf t1 = tf.constant([0, 1, 2, 3, 4, 5, 6, 7]) t1

Output:

<tf.Tensor: shape=(8,), dtype=int32, numpy=array([0, 1, 2, 3, 4, 5, 6, 7], dtype=int32)>

Want to grab a slice from index 1 to 3 (excluding 3)? Use tf.slice:

Python3

s1 = tf.slice(t1, begin=[1], size=[3]) print(s1) # Output: [1 2 3]

Output:

tf.Tensor([1 2 3], shape=(3,), dtype=int32)

Pythonic Slicing:

Feeling more comfortable with Python syntax? Slice directly using square brackets and colons:

Python3

s1 = t1[1:4] # Same output as above

Output:

tf.Tensor([1 2 3], shape=(3,), dtype=int32)

Slicing Rules:

start:stop:step: define the slice range (start included, stop excluded, step between elements).
Omit start or stop to start from the beginning or end.
Negative indices count backwards (e.g., -3: grabs the last 3 elements).

Beyond Single Dimensions:

Slicing works for multi-dimensional tensors too! Imagine a matrix t2 of numbers:

Python3

t2 = tf.constant([[0, 1, 2], [3, 4, 5], [6, 7, 8]]) t2

Output:

<tf.Tensor: shape=(3, 3), dtype=int32, numpy=
array([[0, 1, 2],
       [3, 4, 5],
       [6, 7, 8]], dtype=int32)>

Using tf.slice

Python3

# Get the first row first_row = tf.slice(t2, [0, 0], [1, tf.shape(t2)[1]]) print('First Row:\n', first_row)  # Get the second column second_column = tf.slice(t2, [0, 1], [tf.shape(t2)[0], 1]) print('\nSecond Column:\n', second_column)  # Get the second and third columns second_third_columns = tf.slice(t2, [0, 1], [tf.shape(t2)[0], 2]) print('\nSecond & Third Columns:\n', second_third_columns)

Output:

First Row:
 tf.Tensor([[0 1 2]], shape=(1, 3), dtype=int32)Second Column:
 tf.Tensor(
[[1]
 [4]
 [7]], shape=(3, 1), dtype=int32)Second & Third Columns:
 tf.Tensor(
[[1 2]
 [4 5]
 [7 8]], shape=(3, 2), dtype=int32)

Using Python syntax

Python3

# Get the first row print('First Row:\n', t2[0])  # Get the second column print('\nSecond columns:\n', t2[:, 1])  # Get the second and third columns print('\nSecond & Third columns:\n', t2[:, 1:3])

Output:

First Row:
 tf.Tensor([0 1 2], shape=(3,), dtype=int32)Second columns:
 tf.Tensor([1 4 7], shape=(3,), dtype=int32)Second & Third columns:
 tf.Tensor(
[[1 2]
 [4 5]
 [7 8]], shape=(3, 2), dtype=int32)

Remember, the sliced shape reflects your chosen ranges.

2. Extracting tensors

Extracting in TensorFlow lets you grab specific elements or chunks of your data, like choosing your favorite candies from a mixed bag! Unlike slicing, you don't need them to be neatly lined up - any element or subsection is yours for the taking.

Meet tf.gather:

This handy operation has two inputs:

The treasure chest: your main tensor filled with data.
The treasure map: a list of indices telling you which goodies to grab.

Implementations

Imagine a treasure chest of numbers called t1:

Python3

import tensorflow as tf t1 = tf.constant([0, 1, 2, 3, 4, 5, 6, 7]) t1

Output:

<tf.Tensor: shape=(8,), dtype=int32, numpy=array([0, 1, 2, 3, 4, 5, 6, 7], dtype=int32)>

Want the 0th, 3rd, and 6th elements? Use tf.gather:

Python3

favorites = tf.gather(t1, indices=[0, 3, 6]) print(favorites) # Output: [0 3 6]

Output:

tf.Tensor([0 3 6], shape=(3,), dtype=int32)

Remember, the size of the extracted "loot" matches your treasure map. Here, three elements, so three numbers in the output!

Beyond Flat Lands:

tf.gather works for multi-dimensional treasures too! Imagine a matrix t2 of numbers:

Python3

t2 = tf.constant([[0, 1, 2], [3, 4, 5], [6, 7, 8]]) t2

Output:

<tf.Tensor: shape=(3, 3), dtype=int32, numpy=
array([[0, 1, 2],
       [3, 4, 5],
       [6, 7, 8]], dtype=int32)>

Grab the 3rd row and 1st row, and 3rd column and 1st column:

Python3

Indices = tf.constant([1, 0])  r = tf.gather(t2, indices=Indices, axis=0) print('Along Row\n',r)  c = tf.gather(t2, indices=Indices, axis=1) print('\nAlong Columns\n',c)

Output:

Along Row
 tf.Tensor(
[[3 4 5]
 [0 1 2]], shape=(2, 3), dtype=int32)
Along Columns
 tf.Tensor(
[[1 0]
 [4 3]
 [7 6]], shape=(3, 2), dtype=int32)

The extracted size reflects your choices, like picking two rows of two elements each.

Advanced Picking:

For even more granular control, use tf.gather_nd. It takes:

The treasure chest: same as before.
A multi-layered treasure map: each inner list specifies the location of your desired element/subsection.

Example:

Imagine a 3D treasure chest t3:

Python

t3 = tf.constant([[[1, 3, 5, 7],                    [9, 11, 13, 15]],                   [[17, 19, 21, 23],                    [25, 27, 29, 31]]]) t3

Output:

<tf.Tensor: shape=(2, 2, 4), dtype=int32, numpy=
array([[[ 1,  3,  5,  7],
        [ 9, 11, 13, 15]],
       [[17, 19, 21, 23],
        [25, 27, 29, 31]]], dtype=int32)>

Grab specific elements with tf.gather_nd:

Python3

Indices = [[0, 0, 0], [1, 1, 1], [1, 0, 3]] special_picks = tf.gather_nd(t3, indices=Indices, batch_dims=0) print(special_picks)

Output:

tf.Tensor([ 1 27 23], shape=(3,), dtype=int32)

Explantions:

The value at index 0 ([ 1 ]) is the element at t3[0, 0, 0], which is 1.
The value at index 1 ([ 27 ]) is the element at t3[1, 1, 1], which is 27.
The value at index 2 ([ 23 ]) is the element at t3[1, 0, 3], which is 23.

Python3

print('The value at index t3[0, 0, 0] :',t3[Indices[0]]) print('The value at index t3[1, 1, 1] :',t3[Indices[1]]) print('The value at index t3[1, 0, 3] :',t3[Indices[2]])

Output:

The value at index t3[0, 0, 0] : tf.Tensor(1, shape=(), dtype=int32)
The value at index t3[1, 1, 1] : tf.Tensor(27, shape=(), dtype=int32)
The value at index t3[1, 0, 3] : tf.Tensor(23, shape=(), dtype=int32)

You can also grab subsections using the same principle, with each inner list specifying row and column indices.

3. Inserting tensors

Inserting in TensorFlow lets you plant specific values or chunks of data within your tensors, like carefully placing seeds in a garden! Unlike slicing, there's no need for neat rows - each element or subsection can be placed wherever you choose.

Meet tf.scatter_nd:

This handy tool takes three things:

Planting map: a list of indices telling you where to put the seeds (elements/subsections).
Seed bag: the values you want to insert.
Garden size: the overall shape of your final "planted" tensor.

Example:

Imagine you have an empty garden plot t4 with 8 spots:

Python3

import tensorflow as tf t4 = tf.zeros_like(tf.constant([0, 0, 0, 0, 0, 0, 0, 0])) t4

Output:

<tf.Tensor: shape=(8,), dtype=int32, numpy=array([0, 0, 0, 0, 0, 0, 0, 0], dtype=int32)>

Want to plant seeds with values [10, 20, 30] at spots [1, 3, 6]? Use tf.scatter_nd:

Python3

Index = [[1], [3], [6]] values = [10, 20, 30] t4 = tf.scatter_nd(indices=Index,                     updates=values,                     shape=[8]) print(t4)  # Output: [0 10 0 20 0 0 30 0]

Output:

tf.Tensor([ 0 10  0 20  0  0 30  0], shape=(8,), dtype=int32)

The "garden" grows to match the chosen size, and empty spots stay bare (filled with zeros here).

Beyond Flat Land:

tf.scatter_nd works for multi-dimensional gardens too! Imagine a plot t5 with rows and columns:

Python3

t5 = tf.zeros(shape=(4, 5)) t5

Output:

<tf.Tensor: shape=(4, 5), dtype=float32, numpy=
array([[0., 0., 0., 0., 0.],
       [0., 0., 0., 0., 0.],
       [0., 0., 0., 0., 0.],
       [0., 0., 0., 0., 0.]], dtype=float32)>

Plant seeds [100, 200, 300] at spots [[0, 0], [1, 1], [2, 2]]:

Python3

Index = [[0, 0],           [1, 1],           [2, 2]] values = [100, 200, 300] t5 = tf.scatter_nd(indices=Index,                     updates=values,                     shape=t5.shape) print(t5)  # Output: [[100 0 0 0 0], [0 200 0 0 0], [0 0 300 0 0], [0 0 0 0 0]]

Output:

tf.Tensor(
[[100   0   0   0   0]
 [  0 200   0   0   0]
 [  0   0 300   0   0]
 [  0   0   0   0   0]], shape=(4, 5), dtype=int32)

The "garden" grows to the specified shape, with empty spots remaining bare.

Conclusion

Congratulations! You've embarked on a journey through the world of tensor indexing in TensorFlow, unlocking the ability to access and manipulate data with precision. Let's recap your newfound superpower:

Slicing Like a Pizza Pro:

Use tf.slice to carve out specific chunks of your data - think grabbing a slice of pizza! Define start, end, and step to navigate through single or multiple dimensions.

Picking Your Favorites with tf.gather:

This tool acts like a treasure map, allowing you to extract specific elements or sub-sections based on their indices. Just point and grab!

Multi-Dimensional Treasure Hunting with tf.gather_nd:

Need to delve deeper into complex data structures? This operation lets you pinpoint elements across multiple dimensions with a more advanced treasure map.

Planting Data Seeds with tf.scatter_nd:

Imagine carefully placing seeds in your garden. This operation lets you insert values or sub-sections into specific locations within your tensors, shaping your data as needed.

Bonus: Python Slicing Shorthand:

Use square brackets and colons for quick and convenient slicing within Python syntax, just like slicing a list.

Remember:

The shape of your indexed results depends on the shape of your indices.
Tensor indexing is a key technique for various tasks like cropping, padding, reshaping, and concatenating data.

Now that you've grasped the fundamentals, go forth and explore the limitless possibilities of tensor indexing. Remember, the data is your garden, and you hold the tools to cultivate it into insightful discoveries!

Tensor Indexing in Tensorflow

abhijat_sarari

Improve

Article Tags :

Tensor Indexing in Tensorflow

Tensor in TensorFlow

Characteristics of tensors in TensorFlow

Tensor Indexing in TensorFlow

1. Slicing

Slice It Up with tf.slice:

Pythonic Slicing:

Slicing Rules:

Beyond Single Dimensions:

Using tf.slice

Using Python syntax

2. Extracting tensors

Meet tf.gather:

Implementations

Beyond Flat Lands:

Advanced Picking:

Example:

3. Inserting tensors

Meet tf.scatter_nd:

Example:

Beyond Flat Land:

Conclusion

Slicing Like a Pizza Pro:

Picking Your Favorites with tf.gather:

Multi-Dimensional Treasure Hunting with tf.gather_nd:

Planting Data Seeds with tf.scatter_nd:

Bonus: Python Slicing Shorthand:

Remember:

Similar Reads