Reshape and reduce tensors

Learn it

Words to know

tensor

PyTorch's block of numbers. It can be one number, a list, a table, or have more dimensions.

shape

The size of a tensor along each dimension. A table of 2 rows and 3 columns has shape [2, 3].

dimension

One direction of a tensor. In a table, dimension 0 runs down the rows and dimension 1 across the columns.

How it works

1. torch.tensor(flat, dtype=torch.float64) turns a Python list into a tensor of shape [6].
2. .reshape(rows, cols) reads the same numbers in the same order, filling one row after another.
3. .sum(dim=1) adds along dimension 1, which runs across the columns. Each row collapses to one number. dim=0 would add down the rows and give one number per column.
4. .tolist() turns the result back into a Python list.

Worked example

Input
flat = [1, 2, 3, 4, 5, 6]
rows = 2
cols = 3

Output
[6, 15]

The dimension you name in sum is the one that disappears: shape [2, 3] summed over dimension 1 leaves shape [2].

Why it matters

Every later PyTorch lesson is a chain of shape changes like this one. Reading a shape off each line is the fastest way to find a bug in model code.

Your turn

Using PyTorch, lay a flat list out as a table and add up each row.

flat

a list of rows * cols numbers

rows

the number of rows; 1 or more

cols

the number of columns; 1 or more

Returns

A Python list with the sum of each row, top to bottom.

This lesson needs PyTorch, which does not run in the browser. The app gives you the exercise as a file with its tests, to run on your own computer.

Input
flat = [1, 2, 3, 4, 5, 6]
rows = 2
cols = 3

Output
[6, 15]

Input
flat = [1, 2, 3, 4]
rows = 2
cols = 2

Output
[3, 7]

Input
flat = [-1, 1]
rows = 1
cols = 2

Output
[0]

Hints

Each hint gives away a little more. Open one only when you need it.

1. Hint 1: Where to start

   It is one chain of four calls: make a tensor, reshape it, sum it, turn it back into a list.

2. Hint 2: The key step

   Sum along dimension 1 to get one number per row. Dimension 0 would give one number per column.

3. Hint 3: The outline

   Return torch.tensor(flat, dtype=torch.float64).reshape(rows, cols).sum(dim=1).tolist().

Solution

From here on the page gives the method away. If you are still solving, go back to the hints.

Step by step

1. A tensor of shape [rows * cols]. float64 matches Python's own numbers, so the results agree with the expected values to many digits.
2. The same numbers as a table of shape [rows, cols]. Reshape never changes the order or the count. It fills the first row, then the second, and so on.
3. Add along dimension 1: one total per row. Dimension 1 is the one that is removed. What is left has one entry per row.

The same numbered steps are the comments in the solution.

import torch


def solve(flat, rows, cols):
    # 1. A tensor of shape [rows * cols]
    numbers = torch.tensor(flat, dtype=torch.float64)
    # 2. The same numbers as a table of shape [rows, cols]
    table = numbers.reshape(rows, cols)
    # 3. Add along dimension 1: one total per row
    totals = table.sum(dim=1)
    return totals.tolist()

Watch out

• sum(dim=0) adds down each column and returns cols numbers. For the first example it would give [5, 7, 9].
• Return a Python list. A tensor is not equal to a list, so the checks would fail without .tolist().

Check yourself