One weight matrix, every item

A single weight matrix can transform every item in a batch. There is no need to copy the weights or write a loop:

x:      (batch, time, in_features)
w:      (in_features, out_features)
x @ w:  (batch, time, out_features)

Each row x[b, t] is multiplied by the same w.

Matrix axes and batch axes

When both operands have at least two dimensions, @ treats their last two axes as matrix axes. The left matrix's last dimension must match the right matrix's second-to-last dimension. Any earlier axes are batch axes and must be broadcast-compatible.

(5, 2, 3) @    (3, 4) -> (5, 2, 4)
(1, 2, 3) @ (5, 3, 4) -> (5, 2, 4)

In the first line, one weight matrix is reused for five input matrices. In the second, one input matrix is paired with five weight matrices. The shared size 3 is the dimension summed over, not an output axis.

import torch

x = torch.arange(12, dtype=torch.float32).reshape(2, 3, 2)
w = torch.tensor([[1.0], [2.0]])
y = x @ w
print(list(y.shape))  # [2, 3, 1]
print(y.tolist())  # [[[2.0], [8.0], [14.0]], [[20.0], [26.0], [32.0]]]

Why avoid a loop here?

The starter computes each batch item separately, then stacks the results. One expression lets PyTorch handle the batch and avoids those separate Python calls. This can improve performance, but the benefit depends on shapes, device and execution mode.

Flattening the batch and time axes first is also a valid implementation if you restore the output shape. It is unnecessary for this expression. reshape may share storage or copy depending on strides; adjacent axes alone do not guarantee a view.

What the exercise measures

The harness counts dispatched PyTorch operations, including view operations. For its contiguous inputs, the direct matrix product fits the three-operation limit. That count describes this exercise's execution path, not a promise about every PyTorch version or tensor layout.

Dispatched operations are not the same as GPU kernel launches, elapsed time or allocated bytes. Use a benchmark when choosing an implementation for a real workload.

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.