Chapter 2: Tensor Operations

1. Creating and Initializing Tensors

In Chapter 1, we introduced basic Tensor creation methods such as torch.tensor(), torch.zeros(), torch.ones(), torch.rand(), and torch.randn(). Here, we'll cover additional generation functions that are commonly used in practice, along with how to specify data types (dtype).

Various Generation Functions

PyTorch provides a variety of Tensor generation functions suited to different purposes.

import torch

# Arithmetic sequence (start, end (exclusive), step)
r = torch.arange(0, 10, 2)
print(f"arange: {r}")
# Output: arange: tensor([0, 2, 4, 6, 8])

# Divide the specified range into equal intervals
lin = torch.linspace(0, 1, steps=5)
print(f"linspace: {lin}")
# Output: linspace: tensor([0.0000, 0.2500, 0.5000, 0.7500, 1.0000])

# Identity matrix
identity = torch.eye(3)
print(f"eye:\n{identity}")

# Tensor filled with a specified value
filled = torch.full((2, 3), 7)
print(f"full:\n{filled}")

# Generate a random 3x3 Tensor from integers 0-9
randint_tensor = torch.randint(low=0, high=10, size=(3, 3))
print(f"randint:\n{randint_tensor}")

Specifying and Casting Data Types (dtype)

Every Tensor has a dtype (data type) that describes its elements. If not specified explicitly, it is inferred automatically from the Python numeric literals used (a list of integers becomes int64, and a list containing decimals becomes float32).

import torch

# Create with dtype explicitly specified
int_tensor = torch.tensor([1, 2, 3], dtype=torch.int32)
float_tensor = torch.tensor([1, 2, 3], dtype=torch.float64)
print(f"int32: {int_tensor.dtype}, float64: {float_tensor.dtype}")

# Converting dtype (casting)
x = torch.tensor([1.5, 2.7, 3.9])
print(f"Original dtype: {x.dtype}")

x_int = x.to(torch.int64)   # Convert with .to()
x_int2 = x.long()           # Dedicated methods also work
x_float32 = x_int.float()   # From int to float32

print(f"to(int64): {x_int}, dtype={x_int.dtype}")
print(f".long(): {x_int2}")
print(f".float(): {x_float32}, dtype={x_float32.dtype}")

Creating Tensors from a Template (_like functions)

When you want to create a new Tensor with the same shape, dtype, and device as an existing Tensor, the family of functions ending in _like is convenient.

import torch

original = torch.randn(2, 3, dtype=torch.float64)
print(f"Original Tensor: shape={original.shape}, dtype={original.dtype}")

zeros_copy = torch.zeros_like(original)
ones_copy = torch.ones_like(original)
randn_copy = torch.randn_like(original)

print(f"zeros_like: shape={zeros_copy.shape}, dtype={zeros_copy.dtype}")
print(f"ones_like: shape={ones_copy.shape}, dtype={ones_copy.dtype}")
print(f"randn_like: shape={randn_copy.shape}, dtype={randn_copy.dtype}")

2. Shape Operations (reshape, view, transpose)

In deep learning, you'll frequently need to transform a Tensor's shape — for example, flattening a convolutional layer's output before passing it to a fully connected layer, or rearranging the order of dimensions for batch processing. Here we'll look at the most common shape operations.

The Difference Between reshape() and view()

reshape() and view(), briefly touched on in Chapter 1, both change a Tensor's shape without changing its number of elements, but they behave differently internally. view() can only be used when a Tensor is contiguous in memory, and it shares memory with the original data. reshape() shares memory in the same way as view() whenever possible, but automatically creates a copy to handle the non-contiguous case.

import torch

x = torch.arange(12).reshape(3, 4)
print(f"x:\n{x}")
print(f"is_contiguous: {x.is_contiguous()}")

# Transposing makes the memory layout non-contiguous
x_t = x.t()
print(f"x_t.is_contiguous(): {x_t.is_contiguous()}")

# view() cannot be used on a non-contiguous Tensor
try:
    x_t.view(12)
except RuntimeError as e:
    print(f"Error with view(): {e}")

# reshape() automatically handles this by copying internally
y = x_t.reshape(12)
print(f"reshape() succeeds: {y}")

# If you want to use view(), call contiguous() first
z = x_t.contiguous().view(12)
print(f"contiguous().view() also succeeds: {z}")

transpose() and permute()

transpose(dim0, dim1) swaps the two specified dimensions. When you need to rearrange multiple dimensions at once on a Tensor with three or more dimensions, use permute().

import torch

# Example: convert image data (channel, height, width) to (height, width, channel)
image = torch.randn(3, 32, 32)  # C=3, H=32, W=32
print(f"Original shape (C, H, W): {image.shape}")

# transpose only swaps two dimensions
transposed = image.transpose(0, 2)
print(f"transpose(0, 2): {transposed.shape}")

# permute lets you freely specify the order of dimensions
permuted = image.permute(1, 2, 0)  # (H, W, C)
print(f"permute(1, 2, 0): {permuted.shape}")

Flattening with flatten()

import torch

x = torch.randn(2, 3, 4)

# Flatten the entire tensor to 1D
flat_all = x.flatten()
print(f"flatten(): {flat_all.shape}")  # torch.Size([24])

# Specifying start_dim merges only the dimensions from that point onward
# (a typical use case: flattening while preserving the batch dimension)
flat_from_1 = x.flatten(start_dim=1)
print(f"flatten(start_dim=1): {flat_from_1.shape}")  # torch.Size([2, 12])

3. Indexing and Slicing

If you're familiar with NumPy, its indexing and slicing notation carries over to PyTorch almost unchanged.

Basic Indexing and Slicing

import torch

x = torch.arange(24).reshape(4, 6)
print(f"x:\n{x}")

print(f"x[0]: {x[0]}")               # first row
print(f"x[0, 3]: {x[0, 3]}")         # element at row 1, column 4
print(f"x[1:3]:\n{x[1:3]}")          # rows 2-3
print(f"x[:, 2]: {x[:, 2]}")         # column 3 of all rows
print(f"x[:, 1:4]:\n{x[:, 1:4]}")    # columns 2-4 of all rows
print(f"x[::2]:\n{x[::2]}")          # every other row

# Extract a regular Python number from a scalar Tensor
scalar = x[0, 0]
print(f"scalar: {scalar}, item(): {scalar.item()}")

Indexing with Boolean Masks

To extract only the elements that satisfy a condition, use the result of a comparison operation (a Boolean Tensor) as an index.

import torch

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

mask = x > 0
print(f"mask:\n{mask}")

positive_values = x[mask]
print(f"Positive values only: {positive_values}")

# Rewrite only the elements that satisfy the condition
x[x < 0] = 0
print(f"Negative values replaced with 0:\n{x}")

Advanced Indexing (Fancy Indexing) and Ellipsis

import torch

x = torch.arange(10) * 10  # tensor([0, 10, 20, ..., 90])

# Use an integer Tensor to retrieve multiple elements at once
indices = torch.tensor([0, 2, 5])
print(f"fancy indexing: {x[indices]}")

# For a multi-dimensional Tensor, specify only certain dimensions and omit the rest
y = torch.randn(2, 3, 4, 5)
print(f"y[..., 0].shape: {y[..., 0].shape}")   # select index 0 only along the last dimension
print(f"y[0, ..., 0].shape: {y[0, ..., 0].shape}")

# Advanced indexing returns a copy of the original data (memory is not shared)
sub = x[indices]
sub[0] = -999
print(f"Original x is unchanged: {x[0]}")

4. Mathematical and Matrix Operations

Chapter 1 covered element-wise arithmetic and simple matrix multiplication. Here we'll look at more practical operations, including reduction operations and batch matrix multiplication.

Reduction Operations (sum, mean, max, min, etc.)

import torch

x = torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])

print(f"Sum: {x.sum()}")
print(f"Mean: {x.mean()}")
print(f"Max: {x.max()}")
print(f"Min: {x.min()}")

# The dim argument lets you specify the axis to reduce over
print(f"Sum along columns (dim=0): {x.sum(dim=0)}")   # tensor([5., 7., 9.])
print(f"Sum along rows (dim=1): {x.sum(dim=1)}")   # tensor([6., 15.])

# Get the index of the maximum value
print(f"Index of the overall maximum: {x.argmax()}")
print(f"Index of the maximum per row: {x.argmax(dim=1)}")

# keepdim preserves the number of dimensions while reducing
print(f"keepdim=True: {x.sum(dim=1, keepdim=True).shape}")  # torch.Size([2, 1])

Matrix Multiplication and Batch Matrix Multiplication

import torch

# Matrix multiplication between two 2D Tensors (review from Chapter 1)
A = torch.randn(3, 4)
B = torch.randn(4, 5)
C = A @ B  # Same as torch.matmul(A, B)
print(f"Shape of A @ B: {C.shape}")  # torch.Size([3, 5])

# Batch matrix multiplication: compute multiple matrices at once
batch_A = torch.randn(10, 3, 4)  # batch size 10, each a 3x4 matrix
batch_B = torch.randn(10, 4, 5)  # batch size 10, each a 4x5 matrix

batch_C = torch.bmm(batch_A, batch_B)
print(f"Shape of batch matrix multiplication: {batch_C.shape}")  # torch.Size([10, 3, 5])

# torch.matmul() automatically recognizes batch dimensions, so it can be used instead of bmm()
batch_C2 = torch.matmul(batch_A, batch_B)
print(f"Same result: {torch.allclose(batch_C, batch_C2)}")

Comparison Operations

import torch

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

print(f"x == y: {x == y}")
print(f"x > y: {x > y}")
print(f"torch.eq(x, y): {torch.eq(x, y)}")

# Whether all elements are equal / whether any element satisfies the condition
print(f"All equal: {torch.equal(x, x)}")
print(f"Any True: {(x > 2).any()}")
print(f"All True: {(x > 0).all()}")

5. Broadcasting

Operating on Tensors of different shapes usually raises an error, but when certain rules are satisfied, PyTorch automatically "expands" the shapes to make the operation possible. This is called broadcasting, and it follows the same rules as NumPy's broadcasting.

Broadcasting Rules

When operating on two Tensors, PyTorch compares their shapes using the following procedure.

1. Compare the shapes starting from the trailing (rightmost) dimension
2. Two dimensions are considered compatible if their sizes are equal, or if either one is 1
3. If one Tensor has fewer dimensions, the missing leading dimensions are treated as if padded with size 1
4. If all dimensions are compatible, dimensions of size 1 are virtually "expanded" to match the other Tensor's size (no actual memory is copied)

Concrete Examples

import torch

# Example 1: scalar and Tensor
x = torch.tensor([1.0, 2.0, 3.0])
print(f"x + 10: {x + 10}")  # 10 is added to every element

# Example 2: shape (3,4) and shape (4,)
a = torch.ones(3, 4)
b = torch.tensor([1.0, 2.0, 3.0, 4.0])  # shape (4,)
print(f"Shape of a + b: {(a + b).shape}")  # torch.Size([3, 4])
print(f"a + b:\n{a + b}")

# Example 3: expanding shape (3,1) and shape (1,4) to (3,4)
c = torch.tensor([[1.0], [2.0], [3.0]])       # shape (3, 1)
d = torch.tensor([[10.0, 20.0, 30.0, 40.0]])  # shape (1, 4)
print(f"Shape of c + d: {(c + d).shape}")  # torch.Size([3, 4])
print(f"c + d:\n{c + d}")

# Example 4: 3D and 2D
e = torch.ones(5, 3, 4)
f = torch.randn(3, 4)
print(f"Shape of e + f: {(e + f).shape}")  # torch.Size([5, 3, 4])

When Shapes Are Incompatible

import torch

x = torch.ones(3, 4)
y = torch.ones(3)  # shape (3,) — the trailing dimension does not match x's 4

try:
    z = x + y
except RuntimeError as e:
    print(f"Broadcasting error: {e}")

Practical Example: Feature Normalization

Broadcasting is used extensively when normalizing each feature of batch data using its mean and standard deviation.

import torch

# Data with 5 samples x 3 features
data = torch.tensor([
    [1.0, 100.0, 0.5],
    [2.0, 150.0, 0.7],
    [3.0, 200.0, 0.9],
    [4.0, 250.0, 1.1],
    [5.0, 300.0, 1.3],
])

# Compute the mean and standard deviation for each feature (column)
mean = data.mean(dim=0)   # shape (3,)
std = data.std(dim=0)     # shape (3,)
print(f"Mean: {mean}")
print(f"Standard deviation: {std}")

# Thanks to broadcasting, the (3,) mean can be subtracted from and
# divided across the (5,3) data for every row at once
normalized = (data - mean) / std
print(f"After normalization:\n{normalized}")
print(f"Mean after normalization: {normalized.mean(dim=0)}")  # nearly 0

6. Moving Data Between CPU and GPU

Chapter 1 introduced the basic use of .to(device). Here, we'll look at patterns and caveats commonly encountered in practice.

Specifying a Device and Creating Tensors

import torch

# Automatically detect the available device
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")

# Method 1: move after creation
x = torch.randn(3, 3)
x = x.to(device)

# Method 2: specify the device directly at creation time (efficient, no intermediate copy)
y = torch.randn(3, 3, device=device)

print(f"x.device: {x.device}, y.device: {y.device}")

# You can also move with .cuda() / .cpu() (.cuda() raises an error in environments without a GPU)
if torch.cuda.is_available():
    z = torch.randn(3, 3).cuda()
    print(f"z.device: {z.device}")
    z_back = z.cpu()
    print(f"z_back.device: {z_back.device}")
else:
    print("GPU is not available, so the .cuda() example is skipped")

Errors from Device Mismatches

Operating on Tensors that live on different devices raises an error. Models and data must always be on the same device.

import torch

if torch.cuda.is_available():
    x_cpu = torch.randn(3, 3)
    x_gpu = torch.randn(3, 3).to("cuda")
    try:
        result = x_cpu + x_gpu
    except RuntimeError as e:
        print(f"Device mismatch error: {e}")
    # Correct approach: move one to match the other
    result = x_cpu.to("cuda") + x_gpu
    print(f"Succeeds once devices match: {result.device}")
else:
    print("This environment has no GPU, so the device mismatch error cannot be reproduced.")
    print("Conceptually, directly operating on a CPU Tensor and a GPU Tensor together raises a RuntimeError.")

Moving Multiple Tensors at Once

import torch

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

tensors = {
    "weights": torch.randn(10, 10),
    "bias": torch.randn(10),
    "input": torch.randn(5, 10),
}

# Use a dict comprehension to move everything to the same device at once
tensors = {name: t.to(device) for name, t in tensors.items()}

for name, t in tensors.items():
    print(f"{name}: device={t.device}, shape={t.shape}")

Exercises

import torch

x = torch.randn(2, 3, 4)
permuted = x.permute(2, 0, 1)   # (4, 2, 3)
print(f"After permute: {permuted.shape}")

flattened = permuted.flatten(start_dim=1)  # (4, 6)
print(f"After flatten: {flattened.shape}")

import torch

x = torch.arange(25).reshape(5, 5).float()

# Get the diagonal elements
diagonal = torch.diagonal(x)
print(f"Diagonal elements: {diagonal}")
# Alternative: x[torch.arange(5), torch.arange(5)]

# Subtract the max value of each row (broadcasting)
row_max = x.max(dim=1, keepdim=True).values  # shape (5, 1)
result = x - row_max
print(f"Result after subtracting each row's max:\n{result}")

import torch

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

A = torch.randn(3, 3).to(device)
result = A @ A
result_cpu = result.cpu()

print(f"Computation device: {A.device}")
print(f"Result (on CPU):\n{result_cpu}")

Summary

In this chapter, you learned the fundamentals of freely manipulating PyTorch Tensors.

• ✅ Create and initialize Tensors in various ways, using functions such as arange, linspace, eye, full, and the _like family
• ✅ Freely manipulate Tensor shapes with reshape, view, transpose, permute, and flatten
• ✅ Extract data using basic indexing, slicing, Boolean masks, and fancy indexing
• ✅ Perform practical mathematical and matrix operations, including reduction operations and batch matrix multiplication
• ✅ Understand broadcasting rules and safely perform operations between Tensors of different shapes
• ✅ Move Tensors between CPU and GPU using .to(device)

Reference Resources

• PyTorch Official Documentation: torch.Tensor
• PyTorch Official Documentation: Broadcasting semantics
• PyTorch Official Tutorials
• PyTorch GitHub