🎓 A1Learner
🔍

PyTorch Tensor Shapes

A tensor shape lists how many values sit along each direction of a tensor. The shape (3, 4) means three rows of four values.

Reading a shape

The shape of this table is (2, 3), read aloud as two rows and three columns. It prints as torch.Size([2, 3]).

import torch
table = torch.tensor([[1, 2, 3], [4, 5, 6]])
print(table.shape)

Counting dimensions

Ndim counts the directions of a tensor. For a table, the count is two, the same as its rank.

table = torch.tensor([[1, 2, 3], [4, 5, 6]])
print(table.ndim)

Counting values

Numel counts every value in the tensor. This table holds six values.

table = torch.tensor([[1, 2, 3], [4, 5, 6]])
print(table.numel())

One direction with size()

The size() method answers for one direction. Direction 0 is the rows and direction 1 is the columns.

table = torch.tensor([[1, 2, 3], [4, 5, 6]])
print(table.size(0))
print(table.size(1))

Example

import torch

row = torch.tensor([10, 20, 30, 40])
table = torch.tensor([[1, 2, 3], [4, 5, 6]])
stack = torch.tensor([[[1, 2], [3, 4]], [[5, 6], [7, 8]]])
print(row.shape, row.ndim, row.numel())
print(table.shape, table.ndim, table.numel())
print(stack.shape, stack.ndim, stack.numel())
print(table.size())

Key points

  • A shape tells how long a tensor is in each direction
  • Ndim: the number of directions of a tensor, the same count as its rank
  • Use numel() to count all the values in a tensor
  • Use size() with a direction number to read one direction
  • Direction numbers start at 0, so rows are direction 0

In this example, the program builds a row, a table, and a stack. Each print line shows the shape, the ndim, and the numel of one tensor. The row has shape (4) and one direction. The table has shape (2, 3) and six values. The stack has shape (2, 2, 2) and three directions. The last line prints the table shape again through size().

Note

A printed shape shows as torch.Size([2, 3]) with square brackets.

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

What does the shape (3, 4) describe?