You are given $N$ square matrices $A[0], A[1], \dots, A[N-1]$, each of shape $D \times D$. Define the **inclusive prefix (cumulative) products**: $$Y[i] = A[0] \,@\, A[1] \,@\, \cdots \,@\, A[i]$$ where $@$ denotes ordinary matrix multiplication. Equivalently $Y[0] = A[0]$ and $Y[i] = Y[i-1] \,@\, A[i]$ for $i > 0$. Because matrix multiplication is **associative but not commutative**, the order of factors must be preserved everywhere — in the forward pass, in the gradient, and in any parallel reformulation. Your task is to implement this matrix "cumprod" as you would inside an automatic-differentiation framework: a forward pass producing all $Y[i]$, and a backward pass that, given upstream gradients $G[i] = \partial L / \partial Y[i]$ for every output, returns $\partial L / \partial A[k]$ for every input. You then re-express both passes using a parallel prefix-scan primitive. Treat each part as building on the previous one. Pseudocode or NumPy/PyTorch-style code is fine; correctness of the recurrences and complexity analysis matters more than language details. ### Constraints & Assumptions - Matrices are dense and all the same size $D \times D$; one dense matmul costs $O(D^3)$ time and $O(D^2)$ space. - $N$ can be large (think hundreds to thousands of factors); $D$ is moderate (tens to low hundreds). - Reverse-mode autodiff: a loss $L$ depends on the prefixes, and you are later handed $G[i] = \partial L / \partial Y[i]$ for **every** $i$ (any $G[i]$ may be zero). You must return $\partial L / \partial A[k]$ for every $k$. - "In-place" means $O(1)$ auxiliary space beyond the input/output arrays already given. - Row-major (NumPy/PyTorch) layout; standard reverse-mode conventions. ### Clarifying Questions to Ask - Are the matrices dense and all the same size $D\times D$, or could they be rectangular / sparse / structured? - Do I receive an upstream gradient $G[i]$ for *every* prefix $Y[i]$, or only for the final $Y[N-1]$? (This decides whether backward is one chain or accumulates across all outputs.) - Is the priority minimizing wall-clock latency on a parallel device (favoring the scan), minimizing memory, or minimizing total sequential work? - May I mutate the input array, or must the inputs survive for a later call? - For Part 5, is the scan a true black box (I only get $\oplus$), or may I also save intermediates from the forward scan? ### Part 1 — In-place forward pass and complexity Implement a **simple in-place** forward computation that overwrites the input array so that, on return, slot $i$ holds $Y[i]$, using only $O(1)$ auxiliary storage. State and justify its time and space complexity. ```hint Recurrence You only need the running product so far: $Y[i] = Y[i-1] \,@\, A[i]$. Walk left to right and fold each factor into the accumulated prefix — `A[i] = A[i-1] @ A[i]`. ``` ```hint Cost model Treat one $D\times D$ matmul as $O(D^3)$ and count how many the sweep performs. Then separate the space the *outputs* occupy from the *auxiliary* space the algorithm itself allocates. ``` #### What This Part Should Cover - A correct, order-preserving left-to-right recurrence with the right base case ($Y[0]=A[0]$). - An accurate cost accounting that *separates* output space ($O(ND^2)$, already given) from genuine auxiliary space ($O(1)$ in $N$), and the $O(ND^3)$ time count. - Recognition that the sweep is *lossy* — the original factors are overwritten, foreshadowing Part 2. ### Part 2 — Why in-place breaks autodiff Explain why the in-place version from Part 1 is unsafe for gradient computation in a reverse-mode automatic-differentiation system. Be concrete about *which* intermediate values the backward pass needs and what an autodiff engine does when a saved tensor is mutated. ```hint What backward needs Write the local backward rule for a single product $C = B \,@\, A$. Which *operands* of that node does the chain rule require at backward time, and what did the in-place sweep do to them by the time it finished? ``` ```hint Engine's view Frameworks version-tag tensors and check the version of anything saved for backward. Connect "operand overwritten in place" to two distinct failure modes: a hard versioning error versus a silently wrong gradient read from a stale buffer. ``` #### What This Part Should Cover - Names the specific operands the backward rule consumes — the factor $A[i]$ and the previous prefix $Y[i-1]$ — and shows the in-place sweep destroys exactly those. - Goes beyond "values get overwritten" to the *engine mechanism*: tensor version counters / saved-for-backward checks. - Distinguishes the two outcomes: a loud versioning error vs. a silently incorrect gradient from a mutated buffer. ### Part 3 — Out-of-place forward pass Implement an **out-of-place** forward pass that writes the prefix products to a fresh array $Y$ and preserves everything the backward pass will need. Identify exactly which intermediates must be retained, and state its time and space complexity versus Part 1. ```hint Save list The backward rule for $Y[i] = Y[i-1] \,@\, A[i]$ needs both operands. The original $A[i]$ survive because you did not overwrite them — so which additional array must you keep alive? ``` #### What This Part Should Cover - A fresh-allocation forward that leaves the inputs $A$ intact. - The exact save-list: original factors $A[k]$ plus the exclusive left-prefixes $Y[k-1]$ (with $Y[-1]\equiv I$). - The trade-off vs. Part 1: same $O(ND^3)$ time, but auxiliary memory rises from $O(1)$ to $O(ND^2)$ — the deliberate price of gradient safety. ### Part 4 — Backward pass Given the upstream gradients $G[i] = \partial L / \partial Y[i]$ (one per output, since every prefix may feed the loss), **derive and implement** the backward pass returning $\text{gradA}[k] = \partial L / \partial A[k]$ for all $k$. Give the per-matmul adjoint rules, assemble them into a single reverse recurrence, and state the complexity. Aim for $O(N D^3)$ total work, not $O(N^2 D^3)$. ```hint Matmul adjoints Start from the single node $C = B \,@\, A$ and write the differential $\delta C = \delta B\, A + B\, \delta A$. Pair it against the upstream adjoint $\bar C$ under the Frobenius inner product $\langle X, Y\rangle = \operatorname{tr}(X^{\mathsf T} Y)$ and read off $\bar B$ and $\bar A$ by matching terms — a transpose will appear on the *opposite* operand in each. Then specialize to $Y[i] = Y[i-1] \,@\, A[i]$. ``` ```hint Accumulating across outputs $A[k]$ feeds *every* later prefix $Y[i]$ with $i \ge k$, so its gradient is a sum over consumers. Rather than expand that double sum, carry a single running prefix-adjoint $\bar Y[i]$ that is *seeded* by the directly-supplied $G[i]$ and then *augmented* by whatever flows back from the node that produced $Y[i+1]$. ``` ```hint The recurrence Process the indices in the one order that lets each $\bar Y[i]$ be complete before you use it — which direction is that? At each index you do two things: harvest the contribution to $\text{gradA}[i]$ from the fully-formed $\bar Y[i]$, and propagate the leftover adjoint onto the neighbour it shares a matmul with. Watch the base index where the left operand is the identity rather than a prefix. ``` #### What This Part Should Cover - The two transpose adjoint rules stated correctly: $\bar B = \bar C\,A^{\mathsf T}$ and $\bar A = B^{\mathsf T}\bar C$, derived (not just quoted). - An explanation of why the naive expansion is $O(N^2D^3)$ and how a running prefix-adjoint $\bar Y$ collapses it to a single $O(ND^3)$ sweep. - Correct accumulation: each $A[k]$ gathers gradient from all consumers $Y[i\ge k]$ *and* the chained $Y[i-1]$ dependency, with the $k=0$ / $Y[-1]=I$ base case handled. ### Part 5 — Scan formulation (forward and backward) You are handed a generic **Hillis–Steele inclusive scan** that takes an array and an *associative* binary operator $\oplus$ and returns all inclusive prefixes. (a) Use it to compute the forward prefix products $Y$. (b) Explain how the **backward** pass can also be expressed with scan-style operations. A complete answer gives at least one fully reasoned scan-based backward and relates it to the reverse-mode rules of Part 4. ```hint Forward is one line Matrix multiply is associative, so the inclusive scan with $\oplus(X, Z) = X \,@\, Z$ returns exactly $Y$ — but it is *not* commutative, so keep left/right operands in order. Recall the cost: $O(\log N)$ depth, $O(N\log N)$ total matmuls, versus the sequential $O(N)$. ``` ```hint Backward, route A — differentiate the scan itself A Hillis–Steele scan is nothing but a stack of $O(\log N)$ matmul layers (plus copies for the untouched low indices). So treat it as you would any small network: what do you need to record on the *forward* pass so you can apply the Part-4 matmul adjoints to every combine, and in what order do you revisit the layers? Don't forget the copy nodes also pass gradient through. ``` ```hint Backward, route B — scan an associative composition Look at the gradient you derived in Part 4 and ask whether the per-index adjoint can be written as a recurrence in a single quantity — call it $S[k]$ — running from $k=N-1$ back toward $0$, with $\text{gradA}[k]$ recovered from $S[k]$ and the exclusive left-prefix $Y[k-1]$. If that recurrence has the *affine* shape $S[k] = (\text{constant}_k) + S[k+1] \,@\, (\text{matrix}_k)$, then maps of the form $F(X) = B + X \,@\, M$ are what you are chaining. Compose two such maps symbolically: the result is again one affine map, so the pair $(M, B)$ forms an associative monoid you can scan. Mind the composition order — it is non-commutative, so fold the suffix in the direction the recurrence runs. ``` #### What This Part Should Cover - A one-line forward via associativity, with explicit left/right ordering because $\oplus$ is non-commutative, and the work/depth distinction ($O(N\log N)$ work, $O(\log N)$ depth) vs. the sequential $O(N)$. - At least one *correct* scan-based backward with explicit non-commutativity handling: either differentiating the scan's matmul layers in reverse (with copy-node gradients), or composing affine maps $F(X)=B+X@M$ with the right composition order and fold direction. - A clear tie-back showing the scan backward reproduces the Part-4 gradients. ### What a Strong Answer Covers These dimensions span all five parts; a top answer keeps them consistent throughout rather than re-deriving conventions per part. - **One set of conventions, applied everywhere.** Order-preserving recurrences ($\oplus$ never commuted), the same Frobenius-inner-product adjoint rules in Parts 4–5, and the unifying identity $Y[-1]=I$ used as the exclusive prefix. - **Honest, consistent complexity.** $O(ND^3)$ time for both passes; the memory story ($O(1)$ aux in-place vs. $O(ND^2)$ to retain prefixes); and the scan's work-vs-depth trade ($O(N\log N)$ work, $O(\log N)$ depth) weighed against sequential work. - **Edge cases and stability.** $N=1$, empty input, the $Y[-1]=I$ base case, and the numerical risk of long products (overflow / underflow / ill-conditioning). ### Follow-up Questions - Products of many matrices can overflow, underflow, or become ill-conditioned. How would you detect this and keep both passes stable (periodic rescaling, log-domain tricks, or carrying normalization factors through the gradient)? - A Blelloch (work-efficient) scan does $O(N)$ work in $O(\log N)$ depth versus Hillis–Steele's $O(N\log N)$ work. When would you prefer each here, and how does its backward differ? - In the affine-map scan of route B, each element carries a pair $(M, B)$ of $D\times D$ matrices and composition costs two matmuls. Compare its total work and memory against the sequential backward of Part 4 — when is the parallel version actually worth it? - If the same matrix $A$ is repeated (so $Y[i] = A^{i+1}$), how do the forward and backward simplify, and what is $\partial L/\partial A$? And if gradients flow into only $Y[N-1]$, how do Parts 4 and 5 simplify?