Sharded matrix multiply: forward, backward and planted bugs
Split a matrix multiplication across several devices and explain how the pieces combine. Write the forward pass and the gradient computation for that layout in NumPy or PyTorch. Then a sharded implementation with planted mistakes is handed over: fix it, using a single-device run as the reference for the gradients.
No answer yet. Write the first one; the best answer gets endorsed.