How should I practice Prefix Matrix Products and Autograd?

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

Is Prefix Matrix Products and Autograd associated with OpenAI?

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

Problem · Matrix

Prefix Matrix Products and Autograd

Learn this problem

● Hard

OpenAIFULLTIMEPHONE SCREEN

Problem statement

You are given a sequence of square integer matrices matrices with shape [N, D, D]. Define the inclusive prefix products

P[i] = matrices[0] @ matrices[1] @ ... @ matrices[i],

where @ is ordinary matrix multiplication.

Reverse-mode contract

You are also given upstream with shape [N, D, D]. Treat upstream[i] as the upstream gradient dL / dP[i] of a scalar loss L. Equivalently, the differential of the loss is

dL = sum(i = 0..N-1) <upstream[i], dP[i]>_F,

where <A, B>_F = sum(r, c) A[r][c] * B[r][c]. Compute every prefix product P[i] and every input gradient dL / dmatrices[i].

For the sequential backward step, keep the original input matrices available and use only a constant number of D x D scratch matrices beyond the inputs and returned prefixes and gradients.

Hillis-Steele follow-up

The intended final approach performs the forward pass as an out-of-place Hillis-Steele scan. Start with one snapshot containing the input matrices. For offsets 1, 2, 4, ..., build a new snapshot from the previous one:

If i < offset, copy matrix i unchanged.
Otherwise, set matrix i to previous[i - offset] @ previous[i].

Keep the snapshots needed to reverse these multiplication rounds without mutating an earlier snapshot.

Return format

Return a long[][] with 2 * N rows and D * D columns. Flatten each matrix in row-major order:

Rows 0 through N - 1 contain P[0] through P[N - 1].
Rows N through 2 * N - 1 contain dL / dmatrices[0] through dL / dmatrices[N - 1].

Function

prefixProductAutograd(matrices: int[][][], upstream: int[][][]) → long[][]

Examples

Example 1

matrices = [[[1,2],[3,4]]]upstream = [[[2,0],[1,-1]]]return = [[1,2,3,4],[2,0,1,-1]]

There is one prefix, so P[0] = matrices[0]. Because that prefix is the input itself, its gradient is exactly upstream[0].

Example 2

matrices = [[[1,2],[0,1]],[[2,0],[1,3]],[[1,1],[2,0]]]upstream = [[[1,0],[0,1]],[[0,1],[-1,0]],[[2,-1],[1,1]]]return = [[1,2,0,1],[4,6,1,3],[16,4,7,1],[3,16,2,8],[1,5,3,12],[9,-3,15,-3]]

The first three rows are the row-major forms of P[0], P[1], and P[2]. The final three rows are the corresponding input gradients after contributions from every dependent prefix have been accumulated.

Example 3

matrices = [[[2]],[[-1]],[[3]],[[2]]]upstream = [[[1]],[[2]],[[-1]],[[3]]]return = [[2],[-2],[-6],[-12],[-16],[34],[-10],[-18]]

For D = 1, matrix multiplication becomes scalar multiplication. The four prefixes are 2, -2, -6, and -12; the remaining rows are their accumulated reverse-mode gradients.

Constraints

1 <= N <= 32
1 <= D <= 4
matrices.length = upstream.length = N.
Every matrix in both inputs has exactly D rows and D columns.
Every input entry is an integer from -10 through 10.
Every individual product, partial sum, prefix-product entry, scan-stage entry, reverse-mode adjoint, and final gradient fits in a signed 64-bit integer.
All multiplication and addition are exact; no modulus or floating-point tolerance is used.