Tensor Operations

PyTorch Tensor Operations Part 2: Reduction, Concatenation, Stack and Split

Learning Objectives

After completing this module, you will be able to:

• Understand PyTorch reduction operations and why they are used in deep learning.
• Calculate tensor sum, mean, minimum, maximum, product, variance, and standard deviation.
• Use dim to perform tensor reductions along specific dimensions.
• Find maximum and minimum values with argmax() and argmin().
• Understand the difference between torch.cat() and torch.stack().
• Concatenate PyTorch tensors along existing dimensions.
• Stack tensors along a newly created dimension.
• Split tensors using torch.split() and torch.chunk().
• Understand tensor shapes while combining and dividing data.
• Apply tensor operations to practical machine learning and deep learning workflows.

Introduction to PyTorch Tensor Operations

In the previous module, you learned about arithmetic operations, element-wise operations, broadcasting, matrix multiplication, and dot products.

The next important part of PyTorch tensor operations is learning how to summarize, combine, and divide tensors.

Deep learning models continuously manipulate tensors. For example:

• A model may calculate the average training loss.
• A classification model may find the class with the highest prediction score.
• A computer vision model may concatenate feature maps.
• A training pipeline may stack individual samples into a batch.
• A large tensor may need to be split into smaller tensors.

PyTorch provides optimized functions for these operations.

The four major concepts covered in this module are:

textCopy
1PyTorch Tensor Operations
2│
3├── Reduction Operations
4│   ├── sum
5│   ├── mean
6│   ├── max
7│   ├── min
8│   ├── argmax
9│   ├── argmin
10│   ├── prod
11│   ├── std
12│   └── var
13│
14├── Concatenation
15│   └── torch.cat()
16│
17├── Stacking
18│   └── torch.stack()
19│
20└── Splitting
21    ├── torch.split()
22    └── torch.chunk()

What Are Reduction Operations in PyTorch?

A reduction operation combines multiple tensor values and produces fewer values.

For example:

textCopy
1Input
2
3[2, 4, 6, 8]
4
5      ↓ sum
6
720

Four values have been reduced to one value.

Reduction operations are extremely common in deep learning with PyTorch. They are used for calculating:

• Total values
• Average loss
• Maximum predictions
• Minimum values
• Statistical measurements
• Classification predictions
• Model evaluation metrics

Common PyTorch reduction functions include:

Understanding the dim Parameter

The dim parameter is one of the most important concepts when working with PyTorch tensors.

Consider this matrix:

pythonCopy
1import torch
2
3x = torch.tensor([
4    [1, 2],
5    [3, 4]
6])

Its shape is:

textCopy
1(2, 2)

You can think of the dimensions as:

textCopy
1        dim=1
2      ─────────→
3
4      1   2
5dim=0 3   4
6      ↓

When using a reduction:

pythonCopy
1torch.sum(x, dim=0)

PyTorch reduces dimension 0.

When using:

pythonCopy
1torch.sum(x, dim=1)

PyTorch reduces dimension 1.

A useful rule for a 2D tensor is:

textCopy
1dim=0 → operate vertically → reduce rows → result for each column
2
3dim=1 → operate horizontally → reduce columns → result for each row

Understanding dim becomes especially important when working with neural networks, image tensors, batches, and transformer models.

PyTorch sum() Operation

The torch.sum() function calculates the sum of tensor elements.

Sum of All Elements

pythonCopy
1import torch
2
3x = torch.tensor([1, 2, 3, 4])
4
5result = torch.sum(x)
6
7print(result)

Output:

textCopy
1tensor(10)

Calculation:

textCopy
11 + 2 + 3 + 4 = 10

You can also use the tensor method:

pythonCopy
1print(x.sum())

Both approaches are commonly used.

Sum Along a Dimension

Consider:

pythonCopy
1x = torch.tensor([
2    [1, 2],
3    [3, 4]
4])

Calculate the sum along dim=0:

pythonCopy
1print(torch.sum(x, dim=0))

Output:

textCopy
1tensor([4, 6])

Calculation:

textCopy
1Column 1: 1 + 3 = 4
2
3Column 2: 2 + 4 = 6

Now calculate along dim=1:

pythonCopy
1print(torch.sum(x, dim=1))

Output:

textCopy
1tensor([3, 7])

Calculation:

textCopy
1Row 1: 1 + 2 = 3
2
3Row 2: 3 + 4 = 7

PyTorch mean() Operation

The torch.mean() function calculates the arithmetic average.

For example:

pythonCopy
1import torch
2
3x = torch.tensor([2.0, 4.0, 6.0, 8.0])
4
5print(torch.mean(x))

Output:

textCopy
1tensor(5.)

Calculation:

textCopy
1(2 + 4 + 6 + 8) / 4 = 5

Mean Along a Dimension

pythonCopy
1x = torch.tensor([
2    [1.0, 2.0],
3    [3.0, 4.0]
4])
5
6print(torch.mean(x, dim=0))

Output:

textCopy
1tensor([2., 3.])

The calculation is:

textCopy
1Column 1:
2
3(1 + 3) / 2 = 2
4
5Column 2:
6
7(2 + 4) / 2 = 3

Important Note About mean()

For typical floating-point tensors, torch.mean() works directly.

For integer tensors, convert the tensor to a floating-point dtype when an average is required:

pythonCopy
1x = torch.tensor([1, 2, 3, 4])
2
3print(torch.mean(x.float()))

This is useful because averaging integer values can produce fractional results.

Finding Maximum Values with torch.max()

The torch.max() function finds the largest value.

pythonCopy
1import torch
2
3x = torch.tensor([5, 8, 3, 10])
4
5print(torch.max(x))

Output:

textCopy
1tensor(10)

Maximum Value Along a Dimension

For a 2D tensor:

pythonCopy
1x = torch.tensor([
2    [5, 9],
3    [8, 2]
4])

Use:

pythonCopy
1result = torch.max(x, dim=1)
2
3print(result)

The result contains two pieces of information:

• Maximum values
• Indices where those values occur

Conceptually:

textCopy
1Values:
2
3[9, 8]
4
5Indices:
6
7[1, 0]

For the first row:

textCopy
1[5, 9]
2
3maximum = 9
4index   = 1

For the second row:

textCopy
1[8, 2]
2
3maximum = 8
4index   = 0

You can access them separately:

pythonCopy
1values, indices = torch.max(x, dim=1)
2
3print("Values:", values)
4print("Indices:", indices)

Finding Minimum Values with torch.min()

The torch.min() function returns the smallest value.

pythonCopy
1import torch
2
3x = torch.tensor([4, 2, 9, 1])
4
5print(torch.min(x))

Output:

textCopy
1tensor(1)

You can also find minimum values along a dimension:

pythonCopy
1x = torch.tensor([
2    [4, 2],
3    [9, 1]
4])
5
6values, indices = torch.min(x, dim=1)
7
8print(values)
9print(indices)

Output:

textCopy
1tensor([2, 1])
2tensor([1, 1])

argmax() in PyTorch

torch.argmax() returns the index of the maximum value, not the maximum value itself.

pythonCopy
1import torch
2
3x = torch.tensor([5, 8, 12, 7])
4
5index = torch.argmax(x)
6
7print(index)

Output:

textCopy
1tensor(2)

The tensor is:

textCopy
1Index:  0  1   2  3
2Value:  5  8  12  7
3              ↑
4           maximum

The maximum value is 12, but its index is 2.

This operation is commonly used in PyTorch classification models.

For example:

pythonCopy
1outputs = torch.tensor([1.2, 4.8, 0.7])
2
3predicted_class = torch.argmax(outputs)
4
5print(predicted_class)

Output:

textCopy
1tensor(1)

The model's highest-scoring class is class 1.

argmin() in PyTorch

torch.argmin() returns the index of the smallest value.

pythonCopy
1import torch
2
3x = torch.tensor([5, 8, 12, 1])
4
5index = torch.argmin(x)
6
7print(index)

Output:

textCopy
1tensor(3)

The smallest value is 1, and its index is 3.

Product Reduction with torch.prod()

torch.prod() multiplies all tensor elements.

pythonCopy
1import torch
2
3x = torch.tensor([2, 3, 4])
4
5print(torch.prod(x))

Output:

textCopy
1tensor(24)

Calculation:

textCopy
12 × 3 × 4 = 24

You can also specify a dimension:

pythonCopy
1x = torch.tensor([
2    [1, 2],
3    [3, 4]
4])
5
6print(torch.prod(x, dim=0))

Output:

textCopy
1tensor([3, 8])

Standard Deviation with torch.std()

Standard deviation measures how much values vary around their mean.

pythonCopy
1import torch
2
3x = torch.tensor([2.0, 4.0, 6.0, 8.0])
4
5print(torch.std(x))

The exact result depends on the standard-deviation convention used by the PyTorch version and parameters.

You can explicitly specify the correction when you need a particular statistical definition:

pythonCopy
1print(torch.std(x, correction=0))

For machine learning, understanding standard deviation is useful when working with data normalization, feature scaling, and statistical analysis.

Variance with torch.var()

Variance measures the average squared deviation from the mean.

pythonCopy
1import torch
2
3x = torch.tensor([2.0, 4.0, 6.0, 8.0])
4
5print(torch.var(x))

You can also specify the correction:

pythonCopy
1print(torch.var(x, correction=0))

Variance and standard deviation are frequently encountered in machine learning preprocessing and normalization techniques.

Counting Non-Zero Elements

torch.count_nonzero() counts how many tensor elements are not equal to zero.

pythonCopy
1import torch
2
3x = torch.tensor([0, 1, 5, 0, 8])
4
5print(torch.count_nonzero(x))

Output:

textCopy
1tensor(3)

The non-zero elements are:

textCopy
11, 5, 8

Therefore:

textCopy
1Count = 3

Reduction Operations with a Real Example

Suppose a neural network produces the following loss values:

pythonCopy
1import torch
2
3losses = torch.tensor([0.8, 0.6, 0.4, 0.2])
4
5average_loss = torch.mean(losses)
6
7print(average_loss)

Output:

textCopy
1tensor(0.5000)

This demonstrates a common deep learning pattern:

textCopy
1Individual losses
2       ↓
3[0.8, 0.6, 0.4, 0.2]
4       ↓
5     mean()
6       ↓
7Average loss
8       ↓
9     0.5

What Is Tensor Concatenation?

Tensor concatenation means joining tensors along an existing dimension.

PyTorch provides:

pythonCopy
1torch.cat()

The important idea is:

torch.cat() joins tensors without creating a new dimension.

Concatenating One-Dimensional Tensors

pythonCopy
1import torch
2
3a = torch.tensor([1, 2, 3])
4b = torch.tensor([4, 5, 6])
5
6result = torch.cat((a, b))
7
8print(result)

Output:

textCopy
1tensor([1, 2, 3, 4, 5, 6])

The original shape of each tensor is:

textCopy
1(3,)

The resulting shape is:

textCopy
1(6,)

Concatenating Matrices Along dim=0

Consider:

pythonCopy
1import torch
2
3A = torch.tensor([
4    [1, 2],
5    [3, 4]
6])
7
8B = torch.tensor([
9    [5, 6],
10    [7, 8]
11])

Their shapes are:

textCopy
1A → (2, 2)
2
3B → (2, 2)

Concatenate along dimension 0:

pythonCopy
1C = torch.cat((A, B), dim=0)
2
3print(C)
4print(C.shape)

Output:

textCopy
1tensor([
2    [1, 2],
3    [3, 4],
4    [5, 6],
5    [7, 8]
6])
7
8torch.Size([4, 2])

The number of rows increases:

textCopy
1(2, 2) + (2, 2)
2       ↓
3     (4, 2)

Concatenating Matrices Along dim=1

Now concatenate along dimension 1:

pythonCopy
1C = torch.cat((A, B), dim=1)
2
3print(C)
4print(C.shape)

Output:

textCopy
1tensor([
2    [1, 2, 5, 6],
3    [3, 4, 7, 8]
4])
5
6torch.Size([2, 4])

The number of columns increases:

textCopy
1(2, 2) + (2, 2)
2       ↓
3     (2, 4)

Tensor Concatenation Rules

For torch.cat():

• The tensors must have compatible dimensions.
• Dimensions other than the concatenation dimension must match.
• The selected concatenation dimension can have different sizes.
• The number of dimensions normally remains unchanged.

For example:

textCopy
1(2, 3)
2(4, 3)

Can be concatenated using:

pythonCopy
1torch.cat((A, B), dim=0)

Result:

textCopy
1(6, 3)

But:

textCopy
1(2, 3)
2(4, 2)

cannot be concatenated along dim=0 because their second dimension differs.

What Is Tensor Stacking?

Tensor stacking combines tensors by creating a new dimension.

PyTorch provides:

pythonCopy
1torch.stack()

Unlike torch.cat(), stacking increases the tensor's number of dimensions by one.

Basic torch.stack() Example

pythonCopy
1import torch
2
3a = torch.tensor([1, 2, 3])
4b = torch.tensor([4, 5, 6])
5
6result = torch.stack((a, b))
7
8print(result)
9print(result.shape)

Output:

textCopy
1tensor([
2    [1, 2, 3],
3    [4, 5, 6]
4])
5
6torch.Size([2, 3])

Before stacking:

textCopy
1a → (3,)
2b → (3,)

After stacking:

textCopy
1result → (2, 3)

A new dimension has been created.

Stack Along dim=1

pythonCopy
1result = torch.stack((a, b), dim=1)
2
3print(result)
4print(result.shape)

Output:

textCopy
1tensor([
2    [1, 4],
3    [2, 5],
4    [3, 6]
5])
6
7torch.Size([3, 2])

The position of the new dimension determines the resulting shape.

torch.cat() vs torch.stack()

Understanding the difference between these two functions is essential for PyTorch tensor manipulation.

A simple way to remember:

textCopy
1cat()
2
3Existing dimension
4      ↓
5[1,2] + [3,4]
6      ↓
7[1,2,3,4]

While:

textCopy
1stack()
2
3New dimension
4      ↓
5[1,2]
6[3,4]

Practical Example: Creating a Batch with stack()

Suppose you have individual feature vectors:

pythonCopy
1sample1 = torch.tensor([1.0, 2.0, 3.0])
2sample2 = torch.tensor([4.0, 5.0, 6.0])
3sample3 = torch.tensor([7.0, 8.0, 9.0])

You can create a batch:

pythonCopy
1batch = torch.stack([
2    sample1,
3    sample2,
4    sample3
5])
6
7print(batch)
8print(batch.shape)

Output shape:

textCopy
1torch.Size([3, 3])

The first dimension now represents the number of samples.

This pattern is common in PyTorch DataLoader pipelines and neural network training.

What Is Tensor Splitting?

Tensor splitting is the opposite concept of concatenation.

Instead of joining tensors together, splitting divides one tensor into smaller tensors.

PyTorch provides:

pythonCopy
1torch.split()

and:

pythonCopy
1torch.chunk()

Splitting a Tensor with torch.split()

Consider:

pythonCopy
1import torch
2
3x = torch.tensor([1, 2, 3, 4, 5, 6])

Split it into groups of two:

pythonCopy
1parts = torch.split(x, 2)
2
3print(parts)

Output:

textCopy
1(
2    tensor([1, 2]),
3    tensor([3, 4]),
4    tensor([5, 6])
5)

The result is a tuple containing multiple tensors.

Unequal Tensor Splitting

torch.split() can also accept a list containing the sizes of each part.

pythonCopy
1x = torch.tensor([1, 2, 3, 4, 5, 6])
2
3parts = torch.split(x, [1, 3, 2])
4
5for part in parts:
6    print(part)

Output:

textCopy
1tensor([1])
2
3tensor([2, 3, 4])
4
5tensor([5, 6])

The requested sizes are:

textCopy
11 + 3 + 2 = 6

which matches the number of elements in the original tensor.

Splitting a Matrix

Consider:

pythonCopy
1import torch
2
3A = torch.tensor([
4    [1, 2],
5    [3, 4],
6    [5, 6],
7    [7, 8]
8])

The shape is:

textCopy
1(4, 2)

Split into groups of two rows:

pythonCopy
1parts = torch.split(A, 2, dim=0)
2
3for part in parts:
4    print(part)

Output:

textCopy
1tensor([
2    [1, 2],
3    [3, 4]
4])
5
6tensor([
7    [5, 6],
8    [7, 8]
9])

Understanding torch.chunk()

torch.chunk() divides a tensor into a specified number of chunks.

Example:

pythonCopy
1import torch
2
3x = torch.arange(12)
4
5parts = torch.chunk(x, 3)
6
7for part in parts:
8    print(part)

Output:

textCopy
1tensor([0, 1, 2, 3])
2
3tensor([4, 5, 6, 7])
4
5tensor([8, 9, 10, 11])

There are three chunks.

torch.split() vs torch.chunk()

The difference is mainly how you specify the split.

For example:

pythonCopy
1torch.split(x, 4)

means:

textCopy
1Create pieces containing up to 4 elements each.

Whereas:

pythonCopy
1torch.chunk(x, 3)

means:

textCopy
1Divide the tensor into up to 3 chunks.

When the tensor cannot be divided evenly, the resulting chunks may have different sizes.

Real-World Deep Learning Applications

These operations are not just mathematical exercises. They are used throughout modern PyTorch deep learning workflows.

Calculating Average Loss

pythonCopy
1losses = torch.tensor([0.8, 0.6, 0.4, 0.2])
2
3average_loss = losses.mean()
4
5print(average_loss)

Used when calculating batch-level training statistics.

Selecting a Predicted Class

pythonCopy
1outputs = torch.tensor([
2    [1.2, 4.5, 0.3],
3    [3.8, 1.1, 2.0]
4])
5
6predictions = torch.argmax(outputs, dim=1)
7
8print(predictions)

Output:

textCopy
1tensor([1, 0])

Each row represents one sample, and argmax(dim=1) selects the class with the highest score.

Combining Feature Maps

In computer vision models, tensors may be concatenated along the channel dimension:

pythonCopy
1features = torch.cat((feature_a, feature_b), dim=1)

This type of operation appears in architectures involving feature fusion, skip connections, and convolutional neural networks.

Creating Mini-Batches

Individual samples can be combined using:

pythonCopy
1batch = torch.stack(samples)

This is closely related to how PyTorch DataLoader produces batches for model training.

Splitting Data

A tensor can be divided into smaller pieces:

pythonCopy
1parts = torch.chunk(data, 4)

This can be useful in custom data-processing pipelines and certain parallel-processing workflows.

Common Mistakes with PyTorch Tensor Operations

Mistake 1: Confusing argmax() with max()

This:

pythonCopy
1torch.argmax(x)

returns an index.

This:

pythonCopy
1torch.max(x)

returns the maximum value.

Example:

textCopy
1x = [4, 9, 2]
2
3max     → 9
4argmax → 1

Mistake 2: Using the Wrong Dimension

Consider:

pythonCopy
1x = torch.tensor([
2    [1, 2],
3    [3, 4]
4])

These are different:

pythonCopy
1torch.sum(x, dim=0)

and:

pythonCopy
1torch.sum(x, dim=1)

Always inspect the tensor shape before choosing dim.

Mistake 3: Confusing cat() and stack()

Remember:

textCopy
1cat   → existing dimension
2stack → new dimension

Mistake 4: Ignoring Tensor Shapes

Before concatenation or stacking, check:

pythonCopy
1print(a.shape)
2print(b.shape)

Shape checking is one of the most useful debugging techniques in PyTorch programming.

Best Practices for Tensor Manipulation

Follow these practices when working with PyTorch tensors:

• Check .shape before combining tensors.
• Use dim deliberately instead of relying on defaults.
• Use torch.cat() when extending an existing dimension.
• Use torch.stack() when creating a new dimension.
• Use torch.split() when you know the desired chunk size or sizes.
• Use torch.chunk() when you know the desired number of chunks.
• Use torch.argmax() when you need the position of the highest value.
• Use torch.max() when you need the maximum value.
• Use floating-point tensors when calculating averages or statistical quantities.
• Test tensor operations with small examples before using them inside a large neural network.

Practice Exercise 1: Reduction Operations

Create the following tensor:

pythonCopy
1import torch
2
3x = torch.tensor([
4    [2.0, 4.0, 6.0],
5    [1.0, 3.0, 5.0]
6])

Calculate:

1. Sum of all elements.
2. Mean of all elements.
3. Maximum value.
4. Minimum value.
5. Sum along dim=0.
6. Sum along dim=1.
7. Maximum values along dim=1.
8. Indices of maximum values along dim=1.

Practice Exercise 2: Concatenation

Create:

pythonCopy
1A = torch.tensor([
2    [1, 2],
3    [3, 4]
4])
5
6B = torch.tensor([
7    [5, 6],
8    [7, 8]
9])

Perform:

1. Concatenate along dim=0.
2. Concatenate along dim=1.
3. Print the resulting shapes.
4. Explain why the shapes are different.

Practice Exercise 3: Stack

Create three vectors:

pythonCopy
1a = torch.tensor([1, 2, 3])
2b = torch.tensor([4, 5, 6])
3c = torch.tensor([7, 8, 9])

Perform:

1. Stack along dim=0.
2. Stack along dim=1.
3. Print both shapes.
4. Explain how stacking creates a new dimension.

Practice Exercise 4: Split and Chunk

Create:

pythonCopy
1x = torch.arange(20)

Perform:

1. Split the tensor into groups of five.
2. Split the tensor into four chunks.
3. Print every resulting tensor.
4. Compare torch.split() and torch.chunk().

Knowledge Check

Before moving to the next module, make sure you can answer these questions:

1. What is a reduction operation?
2. What does the dim parameter control?
3. What is the difference between max() and argmax()?
4. What does torch.cat() do?
5. What does torch.stack() do?
6. Why does stack() increase tensor rank?
7. When should you use torch.split()?
8. What is the difference between split() and chunk()?
9. How are reduction operations used when training neural networks?
10. Why is checking tensor shape important in PyTorch?

Module Summary

In this module, you learned how to perform important PyTorch tensor operations used in machine learning and deep learning.

You learned that:

• Reduction operations summarize tensor data.
• torch.sum() calculates totals.
• torch.mean() calculates averages.
• torch.max() and torch.min() find extreme values.
• torch.argmax() and torch.argmin() return the positions of extreme values.
• torch.prod() calculates products.
• torch.std() and torch.var() provide statistical information.
• torch.cat() concatenates tensors along an existing dimension.
• torch.stack() combines tensors along a newly created dimension.
• torch.split() divides tensors using specified sizes.
• torch.chunk() divides tensors into a specified number of chunks.
• Tensor shape and the dim parameter are critical when performing tensor manipulation.

These operations form an important foundation for PyTorch programming, neural networks, computer vision, natural language processing, transformer models, and deep learning.

In the next module, you will learn PyTorch Tensor Reshaping and Shape Manipulation, including reshape(), view(), flatten(), squeeze(), unsqueeze(), transpose(), and permute(). These operations are essential when preparing tensor data for convolutional neural networks, recurrent neural networks, and transformer architectures.