Easy Framework

Broadcasting Addition

Add the same bias vector to every row of a matrix. Each column has its own bias value, reused across all rows.

Your task

Implement add_bias(x, bias) for floating-point inputs with the same dtype and device. The matrix x has shape (m, n) and bias has shape (n,). Return shape (m, n), preserving dtype and device and leaving both inputs unchanged.

output[i, j] = x[i, j] + bias[j]

For x=[[1, 2, 3], [4, 5, 6]] and bias=[10, 20, 30], return [[11, 22, 33], [14, 25, 36]]. Column 0 receives 10 in every row; column 1 receives 20; column 2 receives 30.

Both frameworks handle this with x + bias. Broadcasting aligns dimensions from the right: the vector's length n matches the matrix's last dimension. Its missing row dimension is treated as size 1 and reused across m rows. You do not need to create m copies of the bias.

Keep bias as a vector or a row of shape (1, n). Reshaping it to (n, 1) changes its alignment and can either fail or add along the wrong direction. The diagram shows the intended shapes.

Loading visualization…

Understand the ideas behind this exercise. The Problem tab contains the exact implementation requirements.

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.

Detecting runtime…

Broadcasting Addition

Your task

Match shapes from the right

Check how much memory the result needs

Your code, saved here

Stuck?

JAX reference solution