Medium Framework

Einstein Summation

The trace of a square matrix is the sum of its main diagonal, from top left to bottom right. Compute one trace for each matrix in a batch using einsum.

Your task

Implement batch_trace(x) for a floating-point tensor of shape (batch, n, n). Return a vector of shape (batch,), preserving x's dtype and device and leaving x unchanged. Keep the batch dimension even when there is only one matrix.

output[b] = sum(x[b, i, i] for i from 0 to n-1)

For matrices [[1, 9], [8, 4]] and [[2, 7], [6, 5]], return [5, 7]. The traces are 1 + 4 = 5 and 2 + 5 = 7. Values outside the main diagonal do not contribute, even if they are larger than the diagonal values.

Read the einsum expression

Use torch.einsum('bii->b', x) or jnp.einsum('bii->b', x).

b names the batch axis.
The two occurrences of i in the same input select matching row and column positions: x[b, 0, 0], x[b, 1, 1], and so on.
The output keeps b and omits i, so the diagonal entries are summed within each matrix while the batch entries stay separate.

Repeating a label selects a diagonal; omitting that label from the output requests the sum. These are two distinct parts of the notation. Do not sum all matrix entries, use the opposite diagonal, or reduce the whole batch to one number. The trace uses signed values, not their absolute values.

The diagram below uses the same two matrices as the worked example.

Loading visualization…

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…

Einstein Summation

Your task

Read the einsum expression

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