Tensors, shapes and broadcasting

Problem

A model's input is rarely one number. A grayscale image is a grid of pixels, a sentence is a list of token vectors, and a training step processes 32 of either at once. We need one data structure that can hold all of these, tell us how it is organized, and let arithmetic act on whole grids at once instead of looping over elements in Python.

That structure is the tensor. Three questions come up the moment you use one: what shape is it (so layers fit together), where do its numbers sit in memory (so reshapes are cheap), and what happens when two shapes differ (broadcasting). Get any of them wrong and the code either crashes or, worse, runs and computes the wrong thing.

Intuition

Think of a tensor as nested lists with every level the same length. The number of levels is the rank (also called ndim), and the list of lengths is the shape. A scalar has rank 0 and shape []; a vector of 4 numbers has shape [4]; a 3×4 matrix has shape [3, 4]; a stack of two such matrices has shape [2, 3, 4].

Two ideas carry the rest of the page:

• The numbers are one flat row in memory. The shape is just a way of reading that row. Changing how you read it (a transpose, a reshape, a slice) is usually free; only changing the row itself costs a copy.
• Size-1 axes stretch. When two tensors meet in arithmetic, any axis of length 1 is treated as if it were repeated to match the other side. That is broadcasting, and it is how one bias vector is added to every row of a batch.

Mechanics

Shape, dtype and device

Every tensor carries three facts besides its numbers. Its shape (a torch.Size, a tuple subclass) fixes the layout; len(t) is just t.shape[0] and fails on a rank-0 tensor (Tensor.__len__). Its dtype fixes how each number is stored: float32 by default, int64 for indices, bfloat16 or float16 for fast training, bool for masks. Its device says where the memory lives, cpu or a GPU such as cuda:0. Arithmetic between tensors on different devices is an error, never a silent copy.

Strides, views and copies

A tensor is a pointer into a flat storage plus a shape and a stride per axis: how many elements to skip in storage to move one step along that axis. In the default contiguous (row-major) layout the last axis has stride 1 and each earlier stride is the product of the sizes after it. For shape [2, 3, 4] the strides are [12, 4, 1], so element [1, 2, 3] sits at offset 1·12 + 2·4 + 3·1 = 23 (Tensor.stride).

offset(i₀, …, i_n−1) = Σ_k i_k · stride_k

A view is a new tensor that shares the storage and only changes shape, stride or starting offset. Transposing [2, 3, 4] on axes 0 and 1 gives shape [3, 2, 4] with strides [4, 12, 1]: the strides are swapped, no number moves. The result is no longer contiguous, because walking its last axes no longer walks storage in order (is_contiguous). Writing through a view writes the original.

The cell values are storage offsets computed with the strides shown, so the numbers in each grid are exactly where that view reads from.

Reshape, view, transpose and permute

These operations change the shape and differ in whether they can copy:

• view(shape) never copies. It works only when the new shape can be expressed with strides over the existing storage, roughly: when the axes you merge are already adjacent and in order. Otherwise it raises (Tensor.view).
• reshape(shape) returns a view when it can and a copy when it must, and the docs tell you not to rely on which (torch.reshape). One axis may be -1, inferred from the element count.
• transpose(a, b) swaps two axes and permute(*dims) reorders all of them. Both are always views (transpose, permute).
• unsqueeze(d) inserts a size-1 axis and squeeze(d) removes one; they are how you line up axes before broadcasting (unsqueeze, squeeze).
• contiguous() copies into row-major order if the tensor isn't already, and returns the tensor itself if it is (contiguous). x.transpose(0, 1).contiguous().view(-1) is the classic fix for a view error.

Broadcasting

When an elementwise operation (+, *, comparisons) meets two shapes, PyTorch follows NumPy's rules:

1. Line the shapes up on the right, padding the shorter one with 1s on the left.
2. For each axis, the sizes must be equal or one of them must be 1.
3. The result takes the larger size on each axis; a size-1 side is repeated along it.

The repeat is free: expand makes a view with stride 0 on the stretched axis, so every step along it reads the same element (Tensor.expand, broadcast_to). torch.broadcast_shapes computes the result shape without making any tensor, which is handy for checking a shape before running a model.

The batch dimension

By convention the first axis is the batch: a layer written for one example of shape [features] runs on [batch, features] unchanged, because its weights broadcast across the leading axis. Sequence models add a time axis, [batch, seq, d_model], and images use [batch, channels, height, width] (NCHW). Watch for exceptions: nn.Transformer and the RNN layers default to [seq, batch, …] unless you pass batch_first=True.

Worked example

Implementation

Everything above in a few lines of PyTorch. The comments show what each line prints.

import torch

x = torch.arange(24).reshape(2, 3, 4)
x.shape, x.stride(), x.dtype, x.device   # (torch.Size([2, 3, 4]), (12, 4, 1), torch.int64, device(type='cpu'))

y = x.transpose(0, 1)                    # a view: no copy
y.shape, y.stride()                      # (torch.Size([3, 2, 4]), (4, 12, 1))
y.is_contiguous()                        # False
y.data_ptr() == x.data_ptr()             # True: same storage
y[2, 1, 3] == x[1, 2, 3]                 # tensor(True)
# y.view(-1)                             # RuntimeError: use .reshape(...) or .contiguous().view(...)
flat = y.reshape(-1)                     # copies, because it must

a = torch.tensor([[0], [10], [20]])      # [3, 1]
b = torch.tensor([[0, 1, 2, 3]])         # [1, 4]
(a + b).shape                            # torch.Size([3, 4])
a.expand(3, 4).stride()                  # (1, 0): the stretched axis has stride 0
torch.broadcast_shapes((8, 1, 6, 1), (7, 1, 5))   # torch.Size([8, 7, 6, 5])

acts = torch.randn(32, 10, 512)
bias = torch.zeros(512)
(acts + bias).shape                      # torch.Size([32, 10, 512])
acts.mean(dim=-1, keepdim=True).shape    # torch.Size([32, 10, 1]): ready to broadcast back

And the stride arithmetic from scratch, the same helpers that computed this page's numbers:

def contiguous_strides(shape):
    strides, step = [], 1
    for size in reversed(shape):
        strides.append(step)
        step *= size
    return tuple(reversed(strides))

def broadcast_shape(*shapes):
    rank = max(map(len, shapes))
    padded = [(1,) * (rank - len(s)) + tuple(s) for s in shapes]
    out = []
    for sizes in zip(*padded):
        big = {n for n in sizes if n != 1}
        if len(big) > 1:
            raise ValueError(f"cannot broadcast sizes {sizes}")
        out.append(big.pop() if big else 1)
    return tuple(out)

Tradeoffs

• Views are free but alias. A view costs nothing, but an in-place write to it (y += 1, y[0] = 0) changes the original. Autograd will raise if that write destroys a value it saved for the backward pass.
• Non-contiguous tensors can be slower. Kernels read memory fastest in order. After a permute, an explicit .contiguous() pays one copy to make later ops faster; it is worth it before a heavy op, wasteful before a cheap one.
• Broadcasting hides shape bugs. pred of shape [32, 1] minus target of shape [32] broadcasts to [32, 32], and .mean() of that is a valid-looking, wrong loss. Assert shapes at module boundaries, or keep matching ranks.
• Precision versus memory. float32 is the safe default. bfloat16 halves memory and is faster on modern GPUs at the cost of precision; float64 is rarely worth its cost in deep learning.
• Batch-first versus sequence-first. Batch-first matches how most code and data loaders think; some fused kernels prefer other layouts. Pick one per project and convert at the boundary.

Check yourself

Connections

References

• PyTorch v2.15.0: Tensor.view documentation, including the exact stride condition for a view.
• PyTorch v2.15.0: torch.broadcast_shapes.
• Harris et al. Array programming with NumPy. Nature 585, 2020. doi:10.1038/s41586-020-2649-2. The strided-array model and broadcasting rules PyTorch follows.
• van der Walt, Colbert and Varoquaux. The NumPy Array: A Structure for Efficient Numerical Computation. 2011. arXiv:1102.1523