Implement Matrix Multiplication Forward and Backward (Autograd-Style) in PyTorch

Problem: Implement Matrix Multiplication Forward and Backward (PyTorch) You are given a code skeleton (provided by the interviewer) where you need to implement the forward and backward passes for matrix multiplication. Let: A have shape (M, K) B have shape (K, N) Forward: C = A @ B with shape (M, N) Task 1: Forward Implement forward to return C. Task 2: Backward Given upstream gradient dC of shape (M, N), compute and return: dA = dC @ B^T of shape (M, K) dB = A^T @ dC of shape (K, N) Follow-up Without changing the math, organize/implement the backward in a way inspired by a parallel prefix/scan style (the interviewer referenced something like the Hillis–Steele scan idea). Explain how your approach reduces sequential dependencies or is more parallel-friendly. Constraints Use PyTorch tensor ops (no need to write CUDA). Must pass numerical checks (e.g., match torch.autograd or finite-difference gradients). Ensure correct shape handling and dtype/device consistency.