5 Essential PyTorch Tensor Operations You Need to Master

PyTorch is a Python‑based scientific library that provides advanced operations on tensors—multidimensional arrays of numbers with uniform shape and type. It can run on GPUs and is widely used as a deep‑learning framework.

The five tensor operations covered are:

expand()

permute()

tolist()

narrow()

where()

1. expand()

Expands an existing tensor along dimensions of size 1 to a new size. You can expand any combination of dimensions; set a dimension to -1 to keep its original size.

Note: only a single dimension can be expanded at a time.

# Example 1 - working
a = torch.tensor([[[1,2,3],[4,5,6]]])
a.size()  # torch.Size([1, 2, 3])
a.expand(2,2,3)
# tensor([[[1, 2, 3], [4, 5, 6]],
#         [[1, 2, 3], [4, 5, 6]]])

In this example the original shape [1, 2, 3] is expanded to [2, 2, 3].

2. permute()

Returns a view of the tensor with its dimensions reordered according to the supplied order. For instance, a tensor of shape [1, 2, 3] can be permuted to [3, 2, 1].

# Example 1 - working
a = torch.tensor([[[1,2,3],[4,5,6]]])
a.permute(2,1,0).size()  # torch.Size([3, 2, 1])
a.permute(2,1,0)
# tensor([[[1], [4]],
#         [[2], [5]],
#         [[3], [6]]])

The original dimensions [1, 2, 3] become [3, 2, 1] after permuting, which is useful for re‑ordering tensors before matrix multiplication or other operations.

3. tolist()

Converts the tensor to a Python number, list, or nested list, allowing you to apply any standard Python logic.

# Example 1 - working
a = torch.tensor([[1,2,3],[4,5,6]])
a.tolist()
# [[1, 2, 3], [4, 5, 6]]

The result is a nested Python list representing the tensor's contents.

4. narrow()

Creates a new tensor that is a narrowed view of the original tensor along a specified dimension, starting at a given index and spanning a given length.

# Example 1 - working
a = torch.tensor([[1,2,3,4],[5,6,7,8],[9,10,11,12],[14,15,16,17]])
torch.narrow(a, 1, 2, 2)
# tensor([[ 3,  4],
#         [ 7,  8],
#         [11, 12],
#         [16, 17]])

Here the tensor is narrowed along the second dimension (columns) starting at index 2 and taking two elements, effectively selecting columns 2‑3.

5. where()

Returns a new tensor whose values are chosen from two tensors based on a condition. Elements where the condition is true are taken from the first tensor; otherwise from the second.

# Example 1 - working
a = torch.tensor([[[1,2,3],[4,5,6]]]).to(torch.float32)
b = torch.zeros(1,2,3)
torch.where(a % 2 == 0, b, a)
# tensor([[[1., 0., 3.],
#          [0., 5., 0.]]])

This example replaces even numbers in a with zeros from b, useful for thresholding operations.

Original English article: https://levelup.gitconnected.com/tensor-operations-in-pytorch-798d58e7adb2