Question
Given square matrices `W[0..n-1]`, compute inclusive prefix products `P[i] = W[0] @ ... @ W[i]` using a Hillis–Steele doubling scan, then implement the backward pass without relying on an automatic-differentiation library. ```python def prefix_matmul(weights: np.ndarray) -> np.ndarray: ... def prefix_matmul_backward(weights: np.ndarray, grad_output: np.ndarray) -> np.ndarray: ... ``` ## Requirements - `weights` has shape `(n, d, d)`; return prefixes with the same shape. - Each scan round combines entries `2**r` positions apart, using out-of-place round updates so values from one round are not mixed with the next. - Backward returns the gradient for every input matrix for an arbitrary gradient on every prefix. - Handle `n=0`, `n=1`, non-commutativity, and floating-point dtype consistently. - Explain saved intermediates and complexity in sequential work, parallel depth, and memory. ## Follow-ups 1. Make the forward in-place and identify what backward must save. 2. Map rounds to GPU kernels and reduce synchronization or memory traffic. 3. Support rectangular transforms where compatible.