🎓 A1Learner
🔍

PyTorch Tensor Comparisons

Tensor comparisons test every cell and return True or False for each one. The answers form a new tensor of bool values, with the same shape as the input.

Test every cell

The > operator tests each cell on its own.

scores = torch.tensor([72, 45, 90, 58])
print(scores > 50)

Equal and less than

== tests whether cells are equal. < tests whether cells are smaller.

print(scores == 90)
print(scores < 60)

Count the True cells

True counts as 1 and False counts as 0. Summing a bool tensor counts how many cells passed.

print((scores > 50).sum())

Example

import torch

scores = torch.tensor([[72, 45, 90],
                       [58, 61, 39]])
passing = scores > 50
print(passing)
print(passing.sum())

Key points

  • Comparison: a test applied to every cell, giving one True or False per cell
  • ==: tests whether two cells hold the same value
  • >: tests whether the left cell is bigger
  • <: tests whether it is smaller
  • A comparison result has dtype bool and the same shape as the inputs
  • Use .sum() on a bool tensor to count its True cells, since True counts as 1
  • == on two tensors tests every pair of cells.

In this example, scores is a table of six scores. The test scores > 50 runs on every cell and gives a bool table of the same shape. Four cells hold True. Summing that table counts the True cells, so the count of passing scores is 4.

Note

The result of == is a tensor of answers, not one answer for the whole tensor.

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

scores is a tensor of exam scores. What does scores > 50 return?