Easy Framework

Matrix Multiplication

Multiply two matrices by taking a dot product between each row of A and each column of B.

Your task

Implement matmul(A, B) using the framework's built-in matrix multiplication. A has shape (m, k) and B has shape (k, n). Both are floating-point tensors with the same dtype and device. Return shape (m, n), preserving dtype and device and leaving the inputs unchanged.

output[i, j] = sum(A[i, t] * B[t, j] for t from 0 to k-1)

For A=[[1, 2], [3, 4]] and B=[[5, 6], [7, 8]], return [[19, 22], [43, 50]]. The top-left value is 1*5 + 2*7 = 19: A's first row is paired with B's first column.

The shared inner dimension k is summed over. The remaining dimensions m and n become the output rows and columns. Matrices need not be square: (2, 3) multiplied by (3, 4) produces (2, 4).

Use A @ B, torch.matmul(A, B) or jnp.matmul(A, B). A * B is elementwise multiplication and does not compute these row-column sums. Keep the input order; B @ A may have a different result or incompatible shapes. Keep a matrix result even when m or n is 1.

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

Multiply rows by columns

For A with shape (M, K) and B with shape (K, N), the matrix product has shape (M, N). Each output entry is a dot product of one row of A and one column of B:

$$C_{ij}=\sum_{k=1}^{K}A_{ik}B_{kj}.$$

For example, multiplying row [1, 2] by column [3, 4] gives 1 * 3 + 2 * 4 = 11. Elementwise multiplication would instead give [3, 8]; it does not include the sum.

If examples are rows of X, with shape (B, D), and weights have shape (D, H), then XW produces H features for each of the B examples.

Keep batch axes separate from matrix axes

For batched matrix multiplication, the last two axes hold each matrix. Leading axes identify batches of matrices. A general matrix-multiply operation may broadcast compatible batch axes. A stricter batched operation may require them to match exactly. Use the operation specified by the exercise.

Einstein notation, often used with einsum, names the axes explicitly:

bij,bjk->bik multiplies corresponding matrices in batch b and sums over j.
bi,bj->bij forms an outer product for each example: output (b, i, j) is the product of input (b, i) and input (b, j).

An axis is summed out when its label is absent from the explicit output. A label that appears in both inputs is not necessarily summed: b stays in both outputs above.

Choose a multiplication order with small intermediate results

In exact arithmetic, (AB)C = A(BC), but the two orders can have very different costs. If the shapes are (1000, 10), (10, 10), and (10, 1), forming BC first needs a (10, 1) intermediate. Forming AB first needs a (1000, 10) intermediate and many more operations.

Floating-point rounding can also differ between orders, so do not expect every bit of the two results to match.

Detecting runtime…

Matrix Multiplication

Your task

Multiply rows by columns

Keep batch axes separate from matrix axes

Choose a multiplication order with small intermediate results

Your code, saved here

Stuck?

JAX reference solution