🎓 A1Learner
🔍

PyTorch Slicing Tensors

Slicing takes a run of values from a tensor in one step. A slice uses a start, a stop, and a step for each dimension.

Slice a row of values

A slice names a start, a stop, and a step. The values come out in order.

values = torch.tensor([10, 20, 30, 40, 50, 60])
print(values[1:4])
print(values[::2])

Slice rows and columns

Slicing cuts out a whole run at once. Each dimension gets its own start, stop, and step.

scores = torch.tensor([[85, 92, 78], [90, 88, 95], [72, 81, 84]])
print(scores[:2])
print(scores[:, -1])

Take every second column

A step of 2 skips every second position.

print(scores[:, ::2])

Example

import torch

month = torch.tensor([[1, 2, 3, 4, 5, 6, 7],
                      [8, 9, 10, 11, 12, 13, 14],
                      [15, 16, 17, 18, 19, 20, 21],
                      [22, 23, 24, 25, 26, 27, 28]])

print("Month shape:", month.shape)
print("First two weeks:")
print(month[:2])
print("Last week:")
print(month[3:])
print("Weekdays only:")
print(month[:, :5])

Key points

  • Slice: a run of values cut from a tensor, like the first three rows
  • Use values[1:4] to take the values at positions 1, 2, and 3
  • The stop position is left out, so 1:4 stops before position 4
  • Use table[:2] to take the first two rows of a table
  • Use table[:, -1] to take the last column of a table
  • Leaving out a start or a stop slices from the edge

In this example, month is a table with four rows and seven columns. Each row holds one week of days. month[:2] slices out the first two rows, the first two weeks. month[3:] starts at row 3 and runs to the end. The last slice keeps every row and takes the first five columns. Those five columns are the weekdays.

Note

A colon by itself, as in table[:, -1], means take everything in that dimension

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

What does values[2:5] take from a row of values?