Skip to content

← Broadcasting step 4 of 6

Easy Primitives

What shape comes out

The whole of broadcasting is three rules, applied right to left:

  1. Pad. The shorter shape gets 1s on the left until the ranks match.
  2. Match. Two axes are compatible if they are equal, or if either is 1.
  3. Stretch. The output axis is the larger of the two.

If any pair fails rule 2, the operation raises.

(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)


(3, 4)  and  (3,)                       (2, 3)  and  (3, 2)
(3, 4)                                  (2, 3)
(1, 3)   <- padded                      (3, 2)
------                                  ------
4 vs 3, neither is 1  ->  error         3 vs 2, neither is 1  ->  error

That fourth one is the case from orientation: (3, 4) against (3,) fails, and (4, 4) against (4,) succeeds and gives the wrong answer. Rule 1 is what makes a bare (n,) line up against the last axis rather than the first, which is almost never what someone writing it meant.

Why do it by hand

Because you will spend more time reading broadcasts than writing them, and because the failure mode is a silent wrong answer rather than an exception whenever the two axes happen to agree. Two seconds of padding the shapes on paper is the entire defence.

Your task

def broadcast_shape(a: list[int], b: list[int]) -> list[int] | None

Return the shape the two would broadcast to, or None if they are incompatible. Do not build the tensors.

The starter returns the longer of the two shapes, which is right whenever one shape is a suffix of the other and wrong otherwise.