Skip to content
Medium Framework

Batched Matrix Multiply

Multiply corresponding pairs of matrices in a batch. Pair A[0] with B[0], A[1] with B[1], and so on, without mixing batch entries.

Your task

Implement batched_matmul(A, B) for floating-point tensors with the same dtype and device:

  • A has shape (batch, m, k).
  • B has shape (batch, k, n) with the same batch size.
  • Return shape (batch, m, n), preserving dtype and device. Leave inputs unchanged.
output[s, i, j] = sum(A[s, i, t] * B[s, t, j] for t from 0 to k-1)

Only the inner matrix dimension k is summed over. The leading batch dimension is retained, including when batch=1. No broadcasting between different batch sizes is required in this exercise.

For A=[[[1, 2]], [[3, 4]]] and B=[[[5], [6]], [[7], [8]]], return [[[17]], [[53]]]: the first pair gives 1*5 + 2*6 = 17, and the second gives 3*7 + 4*8 = 53. The output shape is (2, 1, 1), not (2,).

Use torch.bmm(A, B) or A @ B in PyTorch, and jnp.matmul(A, B) or A @ B in JAX. Do not flatten the batch into the matrix dimensions, multiply all possible pairs, or use elementwise A * B.