Katkat
Back

Analyze matrix multiplication complexity

OpenAIMachine Learning Engineer · Technical Screen
HardMachine Learning EngineerTechnical ScreenSoftware Engineering Fundamentals

Last updated

This question evaluates understanding of time and space complexity for matrix multiplication and memory accounting in ML workloads, testing competency in asymptotic analysis, linear algebra operations, and resource estimation.

In an ML coding interview, you're handed a PyTorch file and asked a series of complexity questions about the operations in it. One of them:

Given two dense matrices $A$ and $B$, where $A$ has shape $(m, n)$ and $B$ has shape $(n, p)$, you compute $C = A \mathbin{@} B$ — standard matrix multiplication, as in NumPy/PyTorch.

  1. What is the time complexity of this operation in Big-O notation, in terms of $m$, $n$, and $p$?
  2. What is the space complexity (extra memory usage) of this operation? Clearly state whether you count the output matrix $C$ as part of the space.

Optional follow-up: How does your answer change (if at all) if $A$ and $B$ are batched — e.g. $A$ is $(b, m, n)$, $B$ is $(b, n, p)$, and you compute a batched matmul?

Constraints & Assumptions

  • Matrices are dense (no sparsity to exploit) and stored in standard row-/column-major layout.
  • You're analyzing the standard / library algorithm that A @ B actually runs — not asymptotically-faster sub-cubic algorithms (those can be discussed, but they're not what @ invokes by default).
  • Treat each scalar multiply and add as $O(1)$ work; ignore numerical-precision and overflow concerns for the complexity analysis.
  • "Big-O in terms of $m$, $n$, $p$" — keep the three dimensions distinct; do not collapse to a single variable unless you first state 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 hints the interviewer cares about this distinction explicitly.)
  • Should I assume the naive/standard algorithm, or are you interested in sub-cubic algorithms like Strassen?
  • Is this CPU or GPU, and do you care about the exact FLOP count (the constant factor), or just the asymptotic class?
  • Are the inputs guaranteed dense, or could sparsity change the analysis?

Hints

Hint 1

Where to start: Write out the definition of a single output entry: $C_{ij} = \sum_{k} A_{ik} B_{kj}$. Count the work for one entry, then count how many entries there are.

Hint 2

Time: The matmul is a triple-nested loop over the two output dimensions and the shared **contraction** dimension. Three distinct sizes means a single-letter $O(n^3)$ is wrong — define each variable.

Hint 3

Space: Separate two different quantities: the memory needed to *store the result* $C$ vs. the *auxiliary* scratch space the algorithm needs beyond its inputs and output. Ask yourself: does accumulating $C_{ij}$ require any growing data structure?

Follow-up questions

  • If the matrices were **square** ($m = n = p = N$), what's the complexity, and what's the smallest exponent you know any algorithm achieves in theory vs. in practice?
  • If only **one** operand were batched (e.g. $A$ is $(b, m, n)$ but $B$ is a shared $(n, p)$), can you avoid running $b$ separate matmuls? What's the complexity then?
  • Why does batched matmul (`torch.bmm`) run so 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 instead of dense?