What shape comes out

Broadcasting lets elementwise operations combine compatible shapes. Compare dimensions from the right, padding the shorter shape with leading ones for the comparison.

For each pair of dimensions:

1. If they are equal, keep that size.
2. If one is one, use the other size.
3. Otherwise, the shapes are incompatible.

For example:

(3, 1) and (4,)          (2, 3, 4) and (3, 1)

  (3, 1)                   (2, 3, 4)
  (1, 4)  padded           (1, 3, 1)  padded
  ------                   ---------
  (3, 4)                   (2, 3, 4)

Shapes (3, 4) and (3,) are incompatible: their last dimensions are four and three, and neither is one. (4, 4) and (4,) are compatible, but that does not tell you whether the calculation matches your intent. A vector aligns with the last axis; row weights may need an inserted axis.

Empty dimensions and scalars

The rule is not always to choose the larger size. Zero and one are compatible, and the output size is zero. For example, (0, 3) and (1, 3) broadcast to (0, 3), an empty result. Zero and two are incompatible.

An empty shape list [] represents a scalar, not an empty tensor. For broadcasting, its missing axes act like size-one axes, so it can combine with any valid tensor shape.

import torch

print(list(torch.broadcast_shapes((3, 1), (4,))))  # [3, 4]
print(list(torch.broadcast_shapes((0, 3), (1, 3))))  # [0, 3]
print(list(torch.broadcast_shapes((), (2, 2))))  # [2, 2]

This helper checks shapes without allocating tensors. In this exercise, implement the rules directly so you can trace why each pair succeeds or fails.

Match shapes from the right

Broadcasting lets an operation apply across compatible shapes without explicitly constructing repeated inputs. Compare dimensions from the right: each pair must be equal, or one must be 1. Missing leading dimensions behave like dimensions of 1. At each position, equal sizes stay the same. If one size is 1, the result uses the other size. This also covers empty dimensions: (0, 3) and (1, 3) broadcast to (0, 3), not (1, 3).

For example, (4, 3) plus (3,) applies the same three-feature bias to all four rows. Adding (4, 1) instead applies a different scalar to each row. A size-one axis lets the same value be reused along that axis.

This explains a subtle centring bug: reducing (B, D) over features gives (B,). Subtracting it directly aligns B with the feature axis. If B == D, the operation runs but uses the wrong means. Retaining the reduction axis as (B, 1) expresses the intended dependence.

Check how much memory the result needs

The input need not be physically copied, but the output can still be huge: (N, 1, D) minus (1, M, D) has shape (N, M, D). An in-place operation must fit its result in the existing destination shape. Explicit expansion may also create a view whose positions refer to the same storage, making independent writes unsafe.

Before implementing an operation, write the aligned shapes and the result shape. Then check which inputs are reused and how much space the output needs.