Indexing and slicing tensors in TensorFlow - ML Experiment: Train & Evaluate

Experiment - Indexing and slicing tensors

Problem:You want to extract specific parts of a tensor to use in your model or analysis. Currently, you do not know how to correctly index or slice tensors in TensorFlow.

Current Metrics:No metrics yet, as this is about data manipulation before modeling.

Issue:Without proper indexing and slicing, you cannot efficiently select or manipulate tensor data, which is essential for preparing inputs or inspecting outputs.

Your Task

Learn to extract specific elements, rows, columns, and sub-tensors from a given 3D tensor using TensorFlow indexing and slicing.

Use TensorFlow 2.x API only.

Do not convert tensors to numpy arrays for slicing.

Demonstrate at least four different types of indexing/slicing.

Hint 1

Hint 2

Hint 3

Hint 4

Solution

TensorFlow

import tensorflow as tf

# Create a 3D tensor of shape (3, 4, 5)
tensor = tf.constant([[[i + j + k for k in range(5)] for j in range(4)] for i in range(3)])

# 1. Indexing: Get the first element of the tensor (shape (4,5))
first_element = tensor[0]

# 2. Slicing: Get the first two rows of the first element (shape (2,5))
slice_rows = tensor[0, 0:2, :]

# 3. Indexing: Get the element at position [2,3,4]
single_value = tensor[2, 3, 4]

# 4. Using tf.slice: Extract a sub-tensor starting at [1,1,1] with size [2,2,3]
sub_tensor = tf.slice(tensor, begin=[1, 1, 1], size=[2, 2, 3])

# Print results
print('Original tensor shape:', tensor.shape)
print('First element (tensor[0]) shape:', first_element.shape)
print('Slice rows (tensor[0, 0:2, :]) shape:', slice_rows.shape)
print('Single value (tensor[2, 3, 4]):', single_value.numpy())
print('Sub-tensor (tf.slice):', sub_tensor.numpy())

Created a 3D tensor using tf.constant.

Used direct indexing to get the first element.

Used slicing to get a subset of rows and all columns.

Used tf.slice to extract a sub-tensor with specific start and size.

Results Interpretation

Before: No ability to select parts of tensors, limiting data manipulation.

After: Able to extract specific elements, slices, and sub-tensors efficiently using TensorFlow indexing and slicing.

Mastering indexing and slicing in TensorFlow lets you handle tensor data flexibly, which is crucial for preparing inputs and analyzing outputs in machine learning.

Bonus Experiment

Try to extract a diagonal slice from a 3D tensor and explain how you did it.

💡 Hint

Use tf.linalg.diag_part or combine tf.range with tf.gather_nd for advanced indexing.