← 返回 openai 的题目列表Sharded Matrix Multiplication and Backpropagation
类型：online_judge
Problem: Sharded MatMul and Backpropagation
Given:
Input matrix A with shape m x n
Weight matrix B with shape n x k
Upstream gradient matrix G = dL/dY with shape m x k
The forward pass is:
Y = A @ B
Assume B is column-sharded across p devices:
B = [B_1, B_2, ..., B_p]
Y_i = A @ B_i
Y = concat(Y_1, Y_2, ..., Y_p) along columns
Implement the forward output Y and the backward results:
dA = dL/dA
dB = dL/dB
Mathematically:
dA = G @ B^T
dB = A^T @ G
Your implementation should reflect column sharding:
Compute each local Y_i
Compute each local dB_i
Sum all local contributions to obtain dA
Input Format
m n k p
m rows of A, each with n integers
n rows of B, each with k integers
m rows of G, each with k integers
If k is not divisible by p, earlier shards receive one extra column.
Output Format
Print the three matrices with labels:
Y
...
dA
...
dB
...
Constraints
1 <= m, n, k <= 50
1 <= p <= k
Matrix values are integers in [-100, 100]
Example
Input:
2 2 2 2
1 2
3 4
5 6
7 8
1 0
0 1
Output:
Y
19 22
43 50
dA
5 7
6 8
dB
1 3
2 4
Example
Input
2 2 2 2
1 2
3 4
5 6
7 8
1 0
0 1
Output
Y
19 22
43 50
dA
5 7
6 8
dB
1 3
2 4