How should I practice Column and Row Parallel Linear Layers?

Start by reading the full statement, solve the visible examples, then compare the pattern against related NVIDIA phone screen problems in FastPrep.

Is Column and Row Parallel Linear Layers associated with NVIDIA?

FastPrep tags this as a NVIDIA phone screen coding practice problem when the problem catalog has that company and stage context.

Column and Row Parallel Linear Layers

NVIDIA● MediumFULLTIMEPHONE SCREEN

Learn

Problem statement

For this exercise, simulate the forward pass of an integer linear layer Y = X * W + bias across shards ranks. x is a B × K batch, weight is K × N, and bias has N entries. Ranks have equal contiguous blocks in increasing feature-index order.

mode = "COLUMN": split the output-feature columns of W and the corresponding bias entries. Every rank uses all of X. Its local B × (N / shards) result includes its bias slice. Gather by concatenating feature chunks in rank order for each batch row.
mode = "ROW": split the input-feature columns of X and matching rows of W. Each rank returns a B × N partial product without bias. Sum corresponding entries across ranks, then add the full bias once.

Return a possibly ragged matrix: first all local rows of rank 0, then rank 1, and so on, with batch rows in input order inside each rank; finally append the B complete output rows. Thus there are (shards + 1) * B rows. Column-mode local rows have N / shards entries; row-mode local rows and all final rows have N entries.

The last B rows equal X * W + bias in both modes, but the local rows distinguish the two algorithms. Use ordinary exact integer arithmetic. This exercise models gather and sum mathematically; it requires no GPUs, distributed runtime, activation or backward pass.

Function

parallelLinear(x: int[][], weight: int[][], bias: int[], shards: int, mode: String) → int[][]

Examples

Example 1

x = [[1,2]]weight = [[3,4],[5,6]]bias = [7,8]shards = 2mode = "COLUMN"return = [[20],[24],[20,24]]

Each rank owns one output feature and its bias; gathering [20] and [24] gives [20,24].

Example 2

x = [[1,2]]weight = [[3,4],[5,6]]bias = [7,8]shards = 2mode = "ROW"return = [[3,4],[10,12],[20,24]]

The ranks contribute [3,4] and [10,12] without bias. Sum them, then add [7,8] once.

Example 3

x = [[1,-1],[0,2]]weight = [[2,-3],[4,5]]bias = [1,-2]shards = 1mode = "ROW"return = [[-2,-8],[8,10],[-1,-10],[9,8]]

With one rank, the two local batch rows precede the two final rows. Bias is added only to the final rows.

Constraints

1 <= B <= 4, 1 <= K, N <= 16.
Both input matrices are rectangular and have compatible dimensions; bias.length = N.
1 <= shards <= 4. In column mode, N is divisible by shards; in row mode, K is divisible by shards. No divisibility is required on the other axis.
mode is exactly "COLUMN" or "ROW".
Input and weight entries lie in [-10,10]; bias entries lie in [-100,100].