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? ```hint 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 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 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? ``` ### 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? ### What a Strong Answer Covers - **Defines the variables.** Three distinct dimensions $\Rightarrow$ time $O(mnp)$, not a single-letter $O(n^3)$; reduces to $O(N^3)$ only after explicitly stating $m = n = p = N$. - **Derives time from first principles:** $mp$ output entries, each a length-$n$ dot product $\Rightarrow$ $\Theta(mnp)$, with the bound shown to be tight for the standard algorithm. - **Separates output space from auxiliary space** and names the convention used — $\Theta(mp)$ if counting $C$, $O(1)$ auxiliary otherwise — because that distinction is exactly what part 2 asks for. - **Connects to FLOPs ($\approx 2mnp$)**, the practical reason an ML interviewer asks (compute budgeting / model sizing). - **Knows the boundary of the naive bound:** acknowledges Strassen / sub-cubic algorithms exist but states that NumPy/PyTorch/BLAS run the cubic algorithm with optimized *constants* (cache tiling, SIMD, tensor cores), not a sub-cubic exponent. - **For batching, distinguishes asymptotics from throughput:** time scales linearly in $b$ (no Big-O win) but batched kernels are far faster in wall-clock for hardware-utilization reasons. ### 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?