In-place Matrix Multiplication vs Autograd; Out-of-place Forward/Backward; Scan-based Implementation

AI Coding: Matrix Chain Multiplication, Autograd Safety, Manual Backprop, and Hillis–Steele Scan Given square matrices (A_1, A_2, \dots, A_k) (each (n\times n), floats), compute the matrix-chain product: [ Y = A_1 A_2 \cdots A_k. ] Part 1: In-place version and why it breaks backprop Implement an in-place PyTorch function matmul_chain_inplace(mats) that may reuse/overwrite input storage to save memory, but produces the correct numeric Y. Explain why such in-place/overwrite implementations typically cannot be differentiated correctly by torch.autograd (what autograd assumptions/intermediates get invalidated). Part 2: Out-of-place (autograd-friendly) Implement an out-of-place function matmul_chain(mats) that: does not modify any input tensor in mats; returns Y and supports backward() to obtain gradients for all inputs. Part 3: Manual backward Derive and implement: forward(mats) -> Y backward(mats, dY) -> [dA1, ..., dAk] where dY is the upstream gradient ((n\times n)). Do not rely on autograd for the gradients. Part 4: Hillis–Steele scan for forward/backward Use Hillis–Steele scan (prefix scan) to implement: forward_scan(mats) (prefix products and/or final Y); backward_scan(mats, dY) to compute per-matrix gradients efficiently. State the algebraic property required (e.g., associativity) and provide key code. Constraints (1 \le k \le 10^4) (1 \le n \le 128) Use float32. Example tests k=2, n=2: A1=[[1,2],[3,4]], A2=[[5,6],[7,8]] => Y=[[19,22],[43,50]] Random small case: for k=3,n=3, gradients from backward_scan match torch.autograd within atol=1e-4. Example Input 2 2 1 2 3 4 5 6 7 8 Output 19 22 43 50