🎓 A1Learner
🔍

PyTorch Tensor Reductions

A reduction combines many cells of a tensor into fewer values. sum, mean, max, and min are all reductions.

Add everything up

sum with no dim uses every cell in the tensor.

scores = torch.tensor([[80, 90, 70],
                       [60, 75, 85],
                       [95, 85, 90]])
print(scores.sum())

Mean and min

mean averages the cells. min gives the smallest one.

temps = torch.tensor([20.0, 22.0, 19.0, 25.0])
print(temps.mean())
print(temps.min())

Reduce along one dimension

dim=1 works along each row. dim=0 works down each column.

print(scores.sum(dim=1))
print(scores.sum(dim=0))

Keep the dimension

keepdim=True leaves the reduced dimension at size 1.

totals = scores.sum(dim=1, keepdim=True)
print(totals)
print(totals.shape)

Where the biggest value sits

max returns the largest value. argmax returns the position of that value.

print(scores.max())
print(scores.argmax())

Example

import torch

scores = torch.tensor([[80, 90, 70],
                       [60, 75, 85],
                       [95, 85, 90]])
print(scores.sum(dim=1))
print(scores.sum(dim=0))
print(scores.sum(dim=1).argmax())

Key points

  • Reduction: an operation that combines many cells into fewer values
  • Use sum() to add cells up and mean() to average them
  • Use max() and min() for the largest and smallest values
  • dim: the dimension a reduction runs along, like dim=1 for along each row
  • Use keepdim=True to keep a reduced dimension at size 1
  • argmax: the position of the largest value, not the value itself
  • mean needs float values.

In this example, scores is a table with one row per student and one column per quiz. sum(dim=1) adds along each row, giving one total per student. sum(dim=0) adds down each column, giving one total per quiz. argmax on the student totals gives the position of the biggest total, which is student 2.

Note

On an integer tensor, mean raises an error.

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

A scores table has shape (4, 3), one row per student. What does scores.sum(dim=1) give?