🎓 A1Learner
🔍

PyTorch Boolean Masking

Boolean masking is selection with a bool tensor, called a mask. The mask keeps the values where it is True.

Pick the passing values

A comparison builds the mask. The mask then picks from the scores.

import torch

scores = torch.tensor([45, 80, 62, 90])
mask = scores >= 60
print(mask)
print(scores[mask])

Mask the cells of a table

A mask with the table's shape picks single cells. The result is one flat tensor.

table = torch.tensor([[50, 90], [70, 40]])
print(table[table >= 60])

Keep whole rows

A mask with one value per row keeps or drops each whole row.

table = torch.tensor([[45, 52], [88, 91], [30, 41]])
row_mask = torch.tensor([False, True, False])
print(table[row_mask])

Example

import torch

scores = torch.tensor([[45.0, 52.0],
                       [88.0, 91.0],
                       [67.0, 74.0],
                       [30.0, 41.0]])
averages = scores.mean(dim=1)
passed = averages >= 60
print("Averages:")
print(averages)
print("Passed mask:")
print(passed)
print("Passing rows:")
print(scores[passed])

Key points

  • Mask: a bool tensor that says which values to keep
  • Use a comparison like scores >= 60 to build a mask
  • Use tensor[mask] to collect the values where the mask is True
  • Use a row mask to keep or drop whole rows of a table
  • The mask is a filter, and only the True positions pass through it
  • Masking builds a new tensor and leaves the original unchanged

In this example, scores holds four students with two quiz scores each. The values are decimals so the mean can run. mean(dim=1) averages across each row, giving one average per student. passed compares each average with 60 and builds a row mask. The last print uses the mask on the table itself. Only the two rows marked True in the mask are printed.

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

What does scores[scores >= 60] return?