Introduction to PyTorch Tensor Indexing and Slicing
Tensor indexing and slicing are fundamental PyTorch operations used to access, select, extract, and modify data stored inside tensors.
In real-world deep learning applications, you rarely use an entire tensor for every operation. Instead, you frequently need to select:
- A specific training example
- Individual rows or columns
- Image channels
- A sequence of tokens
- A region of an image
- A batch of samples
- A portion of an embedding
- A specific feature
- A sub-matrix
PyTorch tensor indexing follows many of the same concepts as Python and NumPy indexing, while also supporting multidimensional tensors commonly used in machine learning and deep learning.
Understanding indexing and slicing is therefore essential for working with PyTorch tensors, computer vision, NLP, transformers, CNNs, datasets, and neural networks.
Learning Objectives
After completing this module, you will be able to:
- Understand PyTorch tensor indexing.
- Access individual tensor elements.
- Use positive and negative indexing.
- Index one-dimensional and multidimensional tensors.
- Slice tensors using start, end, and step values.
- Select rows and columns from matrices.
- Extract sub-matrices from tensors.
- Work with three-dimensional tensors.
- Reverse tensors using
torch.flip(). - Select batches, image channels, and sequence data.
- Understand common tensor indexing errors.
- Apply tensor indexing techniques in deep learning workflows.
What Is Tensor Indexing?
Tensor indexing is the process of accessing one or more elements from a PyTorch tensor using their position.
Consider this tensor:
1import torch 2 3x = torch.tensor([10, 20, 30, 40, 50]) 4 5print(x)
Output:
1tensor([10, 20, 30, 40, 50])
Each element has an index:
1Value: 10 20 30 40 50 2Index: 0 1 2 3 4
For example:
1print(x[2])
Output:
1tensor(30)
The index 2 refers to the third element because Python uses zero-based indexing.
Why Tensor Indexing Is Important in Deep Learning
Tensor indexing appears throughout machine learning and deep learning.
For example, a dataset may contain thousands of samples:
1Dataset 2 │ 3 ├── Sample 0 4 ├── Sample 1 5 ├── Sample 2 6 ├── ... 7 └── Sample N
You may need to select only a specific batch:
1batch = data[:32]
For an image tensor, you may need to select a particular channel:
1red_channel = image[0]
For a language model, you may need to select the first few tokens:
1tokens = input_ids[:128]
These are all examples of PyTorch tensor indexing and slicing.
Positive Indexing
Positive indexing starts from 0 at the beginning of the tensor.
Consider:
1import torch 2 3x = torch.tensor([10, 20, 30, 40, 50])
The indexes are:
1Index: 0 1 2 3 4 2Value: 10 20 30 40 50
Access the First Element
1print(x[0])
Output:
1tensor(10)
Access the Third Element
1print(x[2])
Output:
1tensor(30)
Access the Last Element
1print(x[4])
Output:
1tensor(50)
Although x[4] works here, hard-coding the last index is usually less flexible.
A better approach is:
1print(x[-1])
Negative Indexing
Negative indexing accesses elements from the end of the tensor.
For:
1x = torch.tensor([10, 20, 30, 40, 50])
The index mapping is:
1Positive: 0 1 2 3 4 2Value: 10 20 30 40 50 3Negative: -5 -4 -3 -2 -1
Access the Last Element
1print(x[-1])
Output:
1tensor(50)
Access the Second-Last Element
1print(x[-2])
Output:
1tensor(40)
Access the Third-Last Element
1print(x[-3])
Output:
1tensor(30)
Negative indexing is particularly useful when the tensor length is unknown.
Instead of:
1x[len(x) - 1]
you can write:
1x[-1]
Multidimensional Tensor Indexing
PyTorch tensors can contain multiple dimensions.
For a matrix:
1import torch 2 3A = torch.tensor([ 4 [10, 20, 30], 5 [40, 50, 60], 6 [70, 80, 90] 7])
The tensor can be visualized as:
1 Column 2 0 1 2 3 4Row 0 10 20 30 5Row 1 40 50 60 6Row 2 70 80 90
You can access an element using:
1A[row, column]
Access the First Element
1print(A[0, 0])
Output:
1tensor(10)
Access the Center Element
1print(A[1, 1])
Output:
1tensor(50)
Access the Last Element
1print(A[2, 2])
Output:
1tensor(90)
PyTorch also allows:
1A[1][1]
However, the preferred style is generally:
1A[1, 1]
because it expresses multidimensional indexing directly.
Selecting Rows
When you provide only one index for a two-dimensional tensor, PyTorch selects the corresponding row.
1print(A[0])
Output:
1tensor([10, 20, 30])
Select the second row:
1print(A[1])
Output:
1tensor([40, 50, 60])
Select the third row:
1print(A[2])
Output:
1tensor([70, 80, 90])
You can also explicitly select all columns:
1print(A[1, :])
The : means all values along that dimension.
Selecting Columns
To select a column, use : for the row dimension.
First Column
1print(A[:, 0])
Output:
1tensor([10, 40, 70])
Second Column
1print(A[:, 1])
Output:
1tensor([20, 50, 80])
Last Column
1print(A[:, -1])
Output:
1tensor([30, 60, 90])
This pattern is extremely important:
1A[:, column]
means:
1all rows + selected column
What Is Tensor Slicing?
Tensor slicing selects a range of elements from a tensor.
The general Python-style slicing syntax is:
1tensor[start:end:step]
The parameters mean:
| Parameter | Meaning |
|---|---|
start | Starting index |
end | Ending index, excluded |
step | Number of positions to skip |
For example:
1x[1:4]
selects indexes:
11 22 33
but not index 4.
Basic Tensor Slicing
Consider:
1x = torch.tensor([10, 20, 30, 40, 50])
First Three Elements
1print(x[:3])
Output:
1tensor([10, 20, 30])
The omitted start means:
1start = 0
Last Three Elements
1print(x[2:])
Output:
1tensor([30, 40, 50])
Middle Elements
1print(x[1:4])
Output:
1tensor([20, 30, 40])
Entire Tensor
1print(x[:])
Output:
1tensor([10, 20, 30, 40, 50])
Understanding the End Index
One of the most important rules of Python and PyTorch slicing is:
The end index is exclusive.
For example:
1x[1:4]
selects:
1Index 1 → 20 2Index 2 → 30 3Index 3 → 40
It does not include index 4.
Therefore:
1x[1:4] 2 3Start = 1 4End = 4 5 6Selected indexes = 1, 2, 3
Step Slicing
The third value controls the step.
1tensor[start:end:step]
Select Every Second Element
1x = torch.tensor([10, 20, 30, 40, 50, 60]) 2 3print(x[::2])
Output:
1tensor([10, 30, 50])
The indexes selected are:
10 → 10 22 → 30 34 → 50
Select Every Third Element
1print(x[::3])
Output:
1tensor([10, 40])
The indexes selected are:
10 → 10 23 → 40
Slice With Start and Step
1print(x[1:6:2])
Output:
1tensor([20, 40, 60])
The indexes are:
11 → 20 23 → 40 35 → 60
Reversing a PyTorch Tensor
A common Python technique is:
1x[::-1]
However, PyTorch tensor slicing does not support a negative step in this form.
For example:
1x = torch.tensor([10, 20, 30, 40, 50]) 2 3# x[::-1]
Instead, use torch.flip().
1reversed_x = torch.flip(x, dims=[0]) 2 3print(reversed_x)
Output:
1tensor([50, 40, 30, 20, 10])
For a one-dimensional tensor, dims=[0] means that the first dimension is reversed.
Slicing Two-Dimensional Tensors
For a matrix, you can slice both rows and columns.
General syntax:
1tensor[row_start:row_end, column_start:column_end]
Consider:
1A = torch.tensor([ 2 [1, 2, 3], 3 [4, 5, 6], 4 [7, 8, 9] 5])
Select the First Two Rows
1print(A[:2, :])
Output:
1tensor([ 2 [1, 2, 3], 3 [4, 5, 6] 4])
Select the First Two Columns
1print(A[:, :2])
Output:
1tensor([ 2 [1, 2], 3 [4, 5], 4 [7, 8] 5])
Select the First Two Rows and Columns
1print(A[:2, :2])
Output:
1tensor([ 2 [1, 2], 3 [4, 5] 4])
Extracting a Sub-Matrix
You can extract a smaller region from a matrix using two-dimensional slicing.
1A = torch.tensor([ 2 [1, 2, 3], 3 [4, 5, 6], 4 [7, 8, 9] 5]) 6 7submatrix = A[0:2, 1:3] 8 9print(submatrix)
Output:
1tensor([ 2 [2, 3], 3 [5, 6] 4])
The operation means:
1Rows: 20 through 1 3 4Columns: 51 through 2
Visual representation:
1Original: 2 31 2 3 44 5 6 57 8 9 6 7Selected: 8 9 2 3 10 5 6
This type of slicing is useful for extracting regions from image tensors, feature matrices, and numerical datasets.
Three-Dimensional Tensor Indexing
Deep learning models frequently work with tensors containing three or more dimensions.
Create a 3D tensor:
1import torch 2 3x = torch.arange(24).reshape(2, 3, 4) 4 5print(x) 6print(x.shape)
Output:
1torch.Size([2, 3, 4])
The dimensions can be interpreted as:
12 blocks 23 rows per block 34 values per row
Select the First Block
1print(x[0])
This selects the first dimension.
Select a Row From the First Block
1print(x[0, 1])
This selects:
1block = 0 2row = 1
Select a Single Value
1print(x[1, 2, 3])
This selects:
1block = 1 2row = 2 3column = 3
This concept extends to tensors with four, five, or more dimensions.
Tensor Indexing With Ellipsis
For tensors with many dimensions, the ellipsis operator ... can make indexing more readable.
Example:
1x = torch.randn(2, 3, 4, 5)
To select the last element along the final dimension:
1print(x[..., -1])
The ... represents all preceding dimensions.
Conceptually:
1x[..., -1] 2 3Equivalent to: 4 5x[:, :, :, -1]
This technique is particularly useful when working with high-dimensional tensors in neural networks and transformer architectures.
Practical Deep Learning Examples
Tensor indexing becomes especially important when working with batches, images, and sequences.
Select One Sample From a Batch
Suppose a batch of images has shape:
1(batch_size, channels, height, width)
For example:
1images = torch.randn(32, 3, 224, 224)
The shape is:
132 × 3 × 224 × 224
Select the first image:
1image = images[0] 2 3print(image.shape)
Output:
1torch.Size([3, 224, 224])
The batch dimension has been removed because a single sample was selected.
Select the First 8 Images
1batch = images[:8] 2 3print(batch.shape)
Output:
1torch.Size([8, 3, 224, 224])
This is useful when creating smaller mini-batches or debugging model inputs.
Select the First Image Channel
For an image tensor in (C, H, W) format:
1image = images[0] 2 3channel = image[0] 4 5print(channel.shape)
Output:
1torch.Size([224, 224])
The selected channel is now a two-dimensional image plane.
Select a Region of an Image
Suppose:
1image = torch.randn(3, 224, 224)
Extract a region from the first channel:
1patch = image[0, 50:100, 50:100] 2 3print(patch.shape)
Output:
1torch.Size([50, 50])
This technique is useful for image preprocessing, patch extraction, computer vision, and convolutional neural network workflows.
Indexing Sequence Data
Natural language processing models work with token sequences.
Suppose:
1input_ids = torch.tensor([ 2 [101, 2009, 2003, 1037, 3231, 102], 3 [101, 2023, 2003, 2178, 3231, 102] 4])
The shape is:
1(batch_size, sequence_length)
Select the first sequence:
1first_sequence = input_ids[0] 2 3print(first_sequence)
Select the first 3 tokens from every sequence:
1first_tokens = input_ids[:, :3] 2 3print(first_tokens)
This pattern is frequently used in NLP, tokenization, transformer models, attention mechanisms, and language model preprocessing.
Indexing Batch Dimensions
A common convention in deep learning is to keep the batch dimension as the first dimension.
For example:
1(batch, channels, height, width)
A tensor might have:
1(16, 3, 224, 224)
where:
116 = batch size 23 = color channels 3224 = image height 4224 = image width
Selecting:
1images[0]
returns one image.
Selecting:
1images[:8]
returns the first eight images while preserving the batch dimension.
This distinction is important when preparing tensors for neural network models.
Indexing Versus Slicing
Indexing and slicing are related but different.
| Operation | Purpose | Example |
|---|---|---|
| Indexing | Selects a specific position | x[2] |
| Slicing | Selects a range | x[1:4] |
| Column selection | Selects one column | A[:, 1] |
| Row selection | Selects one row | A[1, :] |
| Sub-matrix selection | Selects a region | A[:2, :2] |
A useful mental model is:
1Indexing 2 ↓ 3One specific position 4 5Slicing 6 ↓ 7A range or region
Common Tensor Indexing Errors
Index Out of Range
This code is invalid:
1x = torch.tensor([1, 2, 3]) 2 3print(x[5])
The tensor contains only indexes:
10 21 32
Therefore, index 5 does not exist.
Incorrect Slice Expectations
Remember:
1x[1:4]
does not include index 4.
For:
1x = torch.tensor([10, 20, 30, 40, 50])
the result is:
1tensor([20, 30, 40])
not:
1tensor([20, 30, 40, 50])
Confusing Rows and Columns
For:
1A = torch.tensor([ 2 [1, 2, 3], 3 [4, 5, 6], 4 [7, 8, 9] 5])
This:
1A[1]
selects the second row.
This:
1A[:, 1]
selects the second column.
Remember:
1A[row, column]
Indexing and Autograd
Tensor indexing can also be used with tensors that participate in automatic differentiation.
For example:
1import torch 2 3x = torch.tensor( 4 [1., 2., 3., 4.], 5 requires_grad=True 6) 7 8y = x[1:3] 9 10loss = y.sum() 11 12loss.backward() 13 14print(x.grad)
Output:
1tensor([0., 1., 1., 0.])
Only the selected elements contributed to the sum, so gradients are assigned to those positions.
This illustrates why understanding tensor indexing is important when working with neural network training and PyTorch Autograd.
Best Practices for PyTorch Tensor Indexing
- Use
tensor[row, column]for multidimensional indexing. - Remember that PyTorch uses zero-based indexing.
- Use
-1when you need the last element. - Remember that slice end positions are exclusive.
- Use
torch.flip()when reversing tensors. - Check
.shapebefore performing multidimensional indexing. - Preserve the batch dimension when preparing model inputs.
- Use
...when working with high-dimensional tensors and the preceding dimensions are not important. - Be careful when selecting a single item because indexing can reduce tensor dimensions.
- Use slicing when you need to preserve a dimension while selecting a range.
Practice Exercises
Beginner Exercises
Create:
1import torch 2 3x = torch.tensor([5, 10, 15, 20, 25, 30])
Perform the following operations:
- Access the first element.
- Access the third element.
- Access the last element using negative indexing.
- Select the first four elements.
- Select the last three elements.
- Select every second element.
- Reverse the tensor using
torch.flip().
Matrix Indexing Exercises
Create:
1A = torch.tensor([ 2 [11, 12, 13], 3 [21, 22, 23], 4 [31, 32, 33] 5])
Perform:
- Access
22. - Select the first row.
- Select the last row.
- Select the second column.
- Select the last column.
- Extract the top-left
2 × 2matrix. - Extract the bottom-right
2 × 2matrix. - Extract the first two rows and last two columns.
Deep Learning Exercise
Create a batch of RGB images:
1images = torch.randn(16, 3, 64, 64)
Perform the following:
- Select the first image.
- Select the first 4 images.
- Select the red channel of the first image.
- Extract a
32 × 32patch from the first image. - Select the last image in the batch.
- Print the shape after every operation.
Sequence Data Exercise
Create:
1tokens = torch.arange(40).reshape(4, 10)
The shape represents:
14 sequences 210 tokens per sequence
Perform:
- Select the first sequence.
- Select the first five tokens from every sequence.
- Select the last two tokens from every sequence.
- Select the last sequence.
- Reverse the token order of the first sequence.
Quick Reference
1# Single element 2x[2] 3 4# Last element 5x[-1] 6 7# First three elements 8x[:3] 9 10# Elements from index 2 11x[2:] 12 13# Every second element 14x[::2] 15 16# Reverse 17torch.flip(x, dims=[0]) 18 19# Matrix element 20A[1, 2] 21 22# Row 23A[1, :] 24 25# Column 26A[:, 1] 27 28# Sub-matrix 29A[0:2, 1:3] 30 31# First batch samples 32images[:8] 33 34# First sample 35images[0] 36 37# Ellipsis 38x[..., -1]
Module Summary
In this module, you learned the fundamentals of PyTorch tensor indexing and slicing.
You learned how to:
- Access tensor elements using positive indexing.
- Access elements from the end using negative indexing.
- Index multidimensional tensors using row and column positions.
- Select complete rows and columns.
- Slice one-dimensional tensors.
- Use start, end, and step values.
- Understand why slice end indexes are exclusive.
- Reverse tensors using
torch.flip(). - Extract sub-matrices from two-dimensional tensors.
- Index three-dimensional tensors.
- Use ellipsis with high-dimensional tensors.
- Select images and batches from computer vision datasets.
- Select tokens from sequence and transformer inputs.
- Understand the relationship between tensor indexing and Autograd.
Tensor indexing and slicing are foundational skills for PyTorch deep learning, tensor manipulation, computer vision, natural language processing, CNNs, transformers, image processing, and neural network development.
In the next part, you will move beyond basic positional indexing and learn Boolean Masking, Advanced Indexing, torch.where(), index_select(), gather(), and modifying tensor values. These techniques are especially important for data preprocessing, filtering tensors, classification, attention masks, loss functions, and transformer architectures.