Analyze matrix multiplication complexity

During an ML coding interview, you receive a PyTorch file and are asked several complexity questions about the operations it contains. One question is:

You are given two dense matrices $A$ and $B$, with $A$ of shape $(m, n)$ and $B$ of shape $(n, p)$, and you compute $C = A \mathbin{@} B$, the usual matrix product used by NumPy and PyTorch.

1. In Big-O notation, what is the time complexity of this operation expressed in $m$, $n$, and $p$?
2. What is the space complexity (additional memory) of this operation? Say explicitly whether the output matrix $C$ is included in that space.

Optional follow-up: How, if at all, does your answer change when $A$ and $B$ are batched—for example, $A$ has shape $(b, m, n)$, $B$ has shape $(b, n, p)$, and you perform a batched matmul?

Where to start: Write the formula for one output entry, $C_{ij} = \sum_{k} A_{ik} B_{kj}$. Count the work needed for that single entry, then count the total number of entries.

Time: Matrix multiplication is a triple loop over the two output dimensions and the shared contraction dimension. Because the three sizes are distinct, a one-letter bound like $O(n^3)$ is incorrect; define each variable.

Space: Distinguish two quantities: the memory required to store the result $C$ versus the auxiliary scratch space the algorithm needs beyond its inputs and output. Ask whether accumulating $C_{ij}$ requires any data structure that grows.

Constraints & Assumptions

• The matrices are dense, so there is no sparsity to exploit, and they use standard row- or column-major storage.
• You are analyzing the standard library algorithm that A @ B actually executes, not asymptotically faster sub-cubic methods; those may be mentioned, but they are not what @ calls by default.
• Count each scalar multiplication and addition as $O(1)$ work, and set aside numerical precision and overflow issues for this analysis.
• For Big-O in terms of $m$, $n$, $p$, keep the three dimensions separate; do not combine them into one variable unless you first say the matrices are square.

Clarifying Questions to Ask

• Do you want auxiliary (extra) space only, or total space including the output matrix $C$? Part 2 suggests the interviewer explicitly cares about this distinction.
• Should I assume the naive/standard algorithm, or do you also want sub-cubic approaches such as Strassen?
• Is the target CPU or GPU, and do you care about the exact FLOP count (the constant factor) or only the asymptotic class?
• Are the inputs guaranteed to be dense, or could sparsity alter the analysis?

What a Strong Answer Covers

• Defines the variables. Three distinct dimensions imply time $O(mnp)$, not a one-letter $O(n^3)$; it becomes $O(N^3)$ only after explicitly setting $m = n = p = N$.
• Derives the time from first principles: there are $mp$ output entries, each a dot product of length $n$, giving $\Theta(mnp)$, and shows this bound is tight for the standard algorithm.
• Separates output space from auxiliary space and states the convention: $\Theta(mp)$ when $C$ is counted, otherwise $O(1)$ auxiliary, since that distinction is precisely what part 2 requests.
• Connects the answer to FLOPs, about $2mnp$, which is the practical reason an ML interviewer asks this question: compute budgeting and model sizing.
• Knows the limit of the naive bound: acknowledges that Strassen and other sub-cubic algorithms exist, but states that NumPy, PyTorch, and BLAS run the cubic algorithm with optimized constants such as cache tiling, SIMD, and tensor cores, not a sub-cubic exponent.
• For batching, separates asymptotic behavior from throughput: time grows linearly in $b$, so there is no Big-O improvement, but batched kernels are much faster in wall-clock time because of hardware utilization.

Follow-up Questions

• If the matrices were square, with $m = n = p = N$, what is the complexity, and what is the smallest exponent you know any algorithm achieves in theory versus in practice?
• If only one operand were batched—for example, $A$ has shape $(b, m, n)$ but $B$ is a shared $(n, p)$—can you avoid running $b$ separate matmuls? What would the complexity be then?
• Why does batched matmul (torch.bmm) run much faster than a Python for loop over $b$ separate matmuls, even though both are $O(bmnp)$?
• How would the time complexity change if $A$ were sparse with $\text{nnz}$ nonzeros rather than dense?

Overview: This question tests understanding of time and space complexity for matrix multiplication and memory accounting in ML workloads, covering asymptotic analysis, linear algebra operations, and resource estimation.