How should I practice Tensor-Parallel MLP Gradients?

Start by reading the full statement, solve the visible examples, then compare the pattern against related OpenAI onsite interview problems in FastPrep.

Is Tensor-Parallel MLP Gradients associated with OpenAI?

FastPrep tags this as a OpenAI onsite interview coding practice problem when the problem catalog has that company and stage context.

Problem · Array

Tensor-Parallel MLP Gradients

Learn this problem

● Hard

OpenAIFULLTIMEONSITE INTERVIEW

Problem statement

A two-layer MLP computes hPre = A * x, h = ReLU(hPre), and y = B * h. Both weight matrices are row-sharded equally across numShards. A forward matrix-vector multiplication gathers every shard's output rows. During backpropagation, each shard computes its local contribution to an input gradient, and those contributions must be summed across shards.

Given the complete matrices in shard order, input vector x, upstream gradient yGrad, and shard count, return the correct weight gradients. The result is a ragged matrix containing every row of dA followed by every row of dB.

Use the standard derivatives:

ReLU'(z) = 1 when z > 0, otherwise 0.
For out = W * v, dW is the outer product of dOut and v.
dv = W^T * dOut; with row sharding this is the sum of every shard's local partial vector.

Apply the second-layer backward pass to the post-ReLU activation h, then apply the ReLU gradient using the pre-activation hPre.

Function

shardedMlpGradients(A: double[][], B: double[][], x: double[], yGrad: double[], numShards: int) → double[][]

Examples

Example 1

A = [[1,2],[-1,1]]B = [[2,3],[4,-1]]x = [1,2]yGrad = [1,2]numShards = 2return = [[10,20],[1,2],[5,1],[10,2]]

hPre=[5,1], so both ReLU masks are one. The gathered second-layer input gradient is B^T*yGrad=[10,1]. Thus dA=[[10,20],[1,2]], while dB=outer([1,2],[5,1])=[[5,1],[10,2]].

Example 2

A = [[1,-2],[3,-1]]B = [[1,4]]x = [1,1]yGrad = [3]numShards = 1return = [[0,0],[12,12],[0,6]]

hPre=[-1,2], so the first hidden gradient is masked out. The post-ReLU activation is [0,2], giving dB=[0,6].

Constraints

1 <= A.length, B.length <= 200
1 <= A[0].length <= 200
A.length = B[0].length, x.length = A[0].length, and yGrad.length = B.length.
Both matrix row counts are divisible by numShards.
All rows in each matrix have equal length.
All values are finite and have absolute value at most 100.