Indexing and slicing in PyTorch - Deep Dive

Overview - Indexing and slicing

What is it?

Indexing and slicing in PyTorch means selecting parts of a tensor, like picking pieces from a big cake. Indexing picks single elements or specific positions, while slicing grabs a range or section of the tensor. This helps you work with just the data you need without changing the whole tensor.

Why it matters

Without indexing and slicing, you would have to use the entire tensor every time, which is slow and wasteful. Being able to quickly grab parts of data lets you train models faster and handle complex data easily. It’s like having a sharp knife to cut exactly the piece you want instead of eating the whole cake.

Where it fits

Before learning indexing and slicing, you should understand what tensors are and how they store data. After mastering this, you can learn advanced tensor operations like broadcasting, reshaping, and masking to manipulate data efficiently.

Mental Model

Core Idea

Indexing and slicing let you pick and cut parts of a tensor to focus on specific data without copying or changing the whole thing.

Think of it like...

Imagine a tensor as a big chocolate bar divided into squares. Indexing is like picking one square, and slicing is like breaking off a row or a chunk of squares to eat.

Tensor (3x4):
┌───────────────┐
│ 0  1  2  3    │
│ 4  5  6  7    │
│ 8  9 10 11    │
└───────────────┘

Indexing: tensor[1,2] → 6
Slicing: tensor[0:2, 1:3] → [[1, 2], [5, 6]]

Build-Up - 7 Steps

1

Foundation

Understanding PyTorch tensors basics

Concept: Learn what tensors are and how they store data in multiple dimensions.

A tensor is like a multi-dimensional array. For example, a 2D tensor looks like a matrix with rows and columns. You can create a tensor in PyTorch using torch.tensor(). Example: import torch x = torch.tensor([[1, 2, 3], [4, 5, 6]]) print(x) This prints a 2x3 tensor.

Result

You get a tensor printed as rows and columns, showing the data stored.

Understanding tensors as containers of data in multiple dimensions is key before learning how to pick parts of them.

2

Foundation

Basic indexing to access elements

Concept: Learn how to access single elements in a tensor using indices.

Indexing uses square brackets with positions starting at zero. For a 2D tensor, tensor[row, column] picks one element. Example: x = torch.tensor([[10, 20], [30, 40]]) print(x[0, 1]) # prints 20 Negative indices count from the end, so x[-1, 0] is 30.

Result

You get the single element at the specified position.

Knowing how to pick single elements lets you read or modify specific data points in tensors.

3

Intermediate

Slicing tensors to get sub-tensors

🤔Before reading on: do you think slicing copies data or creates a view? Commit to your answer.

Concept: Slicing extracts a range of elements along each dimension, creating a smaller tensor view.

Use colon ':' to slice ranges. For example, tensor[0:2, 1:3] picks rows 0 and 1, columns 1 and 2. Example: x = torch.tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]]) slice = x[0:2, 1:3] print(slice) This prints: [[2, 3], [5, 6]] Slicing creates a view, not a copy, so changes affect the original tensor.

Result

You get a smaller tensor showing the selected section.

Understanding slicing as creating views helps avoid bugs when modifying data unintentionally.

4

Intermediate

Using advanced indexing with lists and masks

🤔Before reading on: can you use a list of indices to pick non-contiguous elements? Yes or no?

Concept: Advanced indexing lets you pick elements using lists of indices or boolean masks to select complex patterns.

You can pass lists or tensors of indices to pick specific rows or columns. Example: x = torch.tensor([[10, 20], [30, 40], [50, 60]]) rows = torch.tensor([0, 2]) print(x[rows]) # picks rows 0 and 2 Boolean masks select elements where condition is True: mask = x > 25 print(x[mask]) # prints elements greater than 25

Result

You get tensors with only the selected elements or rows.

Advanced indexing allows flexible data selection beyond simple slices, essential for filtering and conditional operations.

5

Intermediate

Indexing with ellipsis and new axes

Concept: Learn how to use ellipsis (...) to simplify indexing and add new dimensions with None.

Ellipsis (...) means 'all remaining dimensions'. For example, tensor[..., 0] picks the first element in the last dimension regardless of tensor shape. Adding a new axis with None or unsqueeze() changes tensor shape: x = torch.tensor([1, 2, 3]) print(x[None, :]) # shape changes from (3,) to (1, 3) This helps prepare tensors for operations needing specific shapes.

Result

You get tensors with selected dimensions or expanded shapes.

Using ellipsis and new axes simplifies working with tensors of varying dimensions and prepares data for broadcasting.

6

Advanced

Memory sharing and copying in slicing

🤔Before reading on: does slicing always create a new copy of data or share memory? Commit to your answer.

Concept: Slicing usually creates a view sharing memory with the original tensor, but some operations force copies.

When you slice a tensor, PyTorch creates a view pointing to the same data. Modifying the slice changes the original tensor. Example: x = torch.tensor([1, 2, 3, 4]) s = x[1:3] s[0] = 100 print(x) # prints [1, 100, 3, 4] To create a copy, use .clone(): c = x[1:3].clone() c[0] = 200 print(x) # original unchanged

Result

You see that slices share data unless explicitly cloned.

Knowing when data is shared prevents bugs where changes unexpectedly affect original tensors.

7

Expert

Indexing performance and advanced tricks

🤔Before reading on: do you think fancy indexing is faster or slower than slicing? Commit to your answer.

Concept: Different indexing methods have different performance costs; understanding this helps optimize code.

Simple slicing is very fast because it creates views without copying. Fancy indexing (using lists or masks) creates copies and is slower. Example: import time x = torch.arange(1000000) start = time.time() y = x[100:200] # slicing print('Slice time:', time.time() - start) start = time.time() y = x[[100, 200, 300]] # fancy indexing print('Fancy index time:', time.time() - start) Use slicing when possible for speed. Also, advanced tricks like using torch.index_select or torch.masked_select can be more efficient for large data.

Result

You observe slicing is faster than fancy indexing.

Understanding performance differences guides writing efficient tensor operations in real projects.

Under the Hood

PyTorch tensors store data in contiguous memory blocks. Indexing with slices creates views by adjusting metadata like shape and strides without copying data. Fancy indexing creates new tensors by copying selected elements. This memory sharing is efficient but requires care to avoid unintended side effects.

Why designed this way?

The design balances speed and flexibility. Views avoid copying for performance, while copies ensure safe independent data. Alternatives like always copying would slow down operations; always sharing would risk bugs. This hybrid approach is a practical tradeoff.

Tensor data memory:
┌─────────────────────────────┐
│ [data block in memory]      │
│  1  2  3  4  5  6  7  8    │
└─────────────────────────────┘

Slicing view:
┌─────────────┐
│ shape, strides │
│ points to data │
└─────────────┘

Fancy indexing:
┌─────────────┐
│ new data copy │
└─────────────┘

Myth Busters - 4 Common Misconceptions

Quick: Does slicing always create a new copy of the tensor data? Commit yes or no.

Common Belief: Slicing always makes a new copy, so changing a slice won't affect the original tensor.

Reality: Slicing creates a view sharing the same memory, so modifying the slice changes the original tensor.

Why it matters: Assuming slicing copies data can lead to bugs where changes unexpectedly affect the original tensor, causing hard-to-find errors.

Quick: Can you use negative indices in PyTorch indexing? Commit yes or no.

Common Belief: Negative indices are not allowed in PyTorch indexing.

Reality: PyTorch supports negative indices to count from the end of a dimension.

Why it matters: Not knowing this limits your ability to write concise and flexible code for accessing tensor elements.

Quick: Does fancy indexing return a view or a copy? Commit your answer.

Common Belief: Fancy indexing returns a view like slicing.

Reality: Fancy indexing returns a new tensor copy, not a view.

Why it matters: Misunderstanding this can cause performance issues and unexpected behavior when modifying tensors.

Quick: Is ellipsis (...) required to index all dimensions? Commit yes or no.

Common Belief: You must specify all dimensions explicitly; ellipsis is just a shortcut with no real effect.

Reality: Ellipsis lets you skip specifying some dimensions, making code cleaner and more flexible.

Why it matters: Ignoring ellipsis leads to verbose and less maintainable code, especially with high-dimensional tensors.

Expert Zone

1

Slicing creates views that share memory, but some operations like transpose or permute may create non-contiguous tensors requiring copying for some operations.

2

Fancy indexing can trigger unexpected data copies and slowdowns, so profiling code to choose between slicing, index_select, or masked_select is crucial in performance-critical applications.

3

Using None (or unsqueeze) to add dimensions is essential for broadcasting, but forgetting to do so can cause silent shape mismatches or inefficient computations.

When NOT to use

Avoid fancy indexing when working with very large tensors or in tight loops due to copying overhead; instead, use slicing or torch.index_select. For complex conditional selection, consider boolean masks or torch.where. When you need independent data, explicitly clone slices to avoid side effects.

Production Patterns

In production, slicing is used for batching data and extracting features efficiently. Fancy indexing is common in data augmentation and filtering. Experts combine indexing with broadcasting and masking to write concise, fast data pipelines and model inputs.

Connections

Broadcasting

Indexing with new axes (None) prepares tensors for broadcasting by aligning dimensions.

Understanding indexing helps you reshape tensors correctly to leverage broadcasting for efficient arithmetic operations.

Database querying

Indexing and slicing in tensors is similar to selecting rows and columns in database tables.

Knowing how databases filter and select data helps grasp tensor indexing as a way to query multi-dimensional data efficiently.

Photography cropping

Slicing a tensor is like cropping a photo to focus on a part of the image.

This connection shows how selecting parts of data is a universal concept across fields, helping understand the purpose of slicing.

Common Pitfalls

#1 Modifying a slice expecting original tensor to stay unchanged.

Wrong approach:

x = torch.tensor([1, 2, 3, 4])
s = x[1:3]
s[0] = 100
print(x) # expects [1, 2, 3, 4], but prints [1, 100, 3, 4]

Correct approach:

x = torch.tensor([1, 2, 3, 4])
s = x[1:3].clone()
s[0] = 100
print(x) # prints [1, 2, 3, 4]

Root cause:

Not realizing slicing creates a view sharing memory, so changes affect the original tensor.

#2 Using fancy indexing inside a loop causing slowdowns.

Wrong approach:

for i in range(1000):
y = x[[i, i+1]] # fancy indexing inside loop

Correct approach:

indices = torch.arange(0, 1001)
y = x[indices] # use slicing or batch indexing outside loop

Root cause:

Not understanding that fancy indexing copies data and is slower than slicing.

#3 Forgetting to add new axis when needed for broadcasting.

Wrong approach:

x = torch.tensor([1, 2, 3])
y = torch.tensor([[10], [20], [30]])
result = x + y # shape mismatch error

Correct approach:

x = torch.tensor([1, 2, 3])
x = x[None, :] # add new axis
result = x + y # works correctly

Root cause:

Not using None or unsqueeze to align tensor shapes for broadcasting.

Key Takeaways

Indexing and slicing let you select parts of tensors efficiently without copying data unnecessarily.

Slicing creates views sharing memory, so modifying slices affects the original tensor unless you clone explicitly.

Fancy indexing returns copies and is slower, so use it carefully in performance-critical code.

Using ellipsis and new axes simplifies working with tensors of different shapes and prepares data for broadcasting.

Understanding these concepts deeply helps write faster, cleaner, and bug-free tensor manipulation code in PyTorch.