Katkat

Implement Prefix Products and Their Backward Pass

OpenAIMachine Learning Engineer · Onsite

Last updated

Work through inclusive prefix products in separate-output, in-place, reverse-gradient, and parallel variants. Explain zeros, overflow, mutation and aliasing, constant temporary workspace, synchronization, floating-point effects, and the work/depth trade-offs of each part.

Let x[0..n-1] be a sequence and define inclusive prefix products by y[i] = x[0] * x[1] * ... * x[i]. Work through the following variants. State how zeros and numeric overflow are handled by your chosen numeric type.

Part 1: Out-of-Place Forward Pass

Compute all prefix products into a separate output array.

What This Part Should Cover

  • The running-product invariant
  • Linear work and output-space accounting
  • Empty-input behavior

Part 2: In-Place Forward Pass

Overwrite the input with its prefix products without using another array.

What This Part Should Cover

  • Safe update order
  • Why the previous original value is no longer needed
  • Mutation and aliasing assumptions

Part 3: Backward Pass

Given upstream gradients for every prefix output, derive gradients for every original input. Give a reverse-mode algorithm using constant scalar workspace beyond the required input, output, and gradient arrays. Do not rely on division, so zeros remain valid.

What This Part Should Cover

  • The local recurrence y[i] = y[i-1] * x[i]
  • Reverse accumulation into both parents
  • Zero-safe reasoning

Part 4: Hillis-Steele Scan

Describe a parallel Hillis-Steele forward scan and its backward pass. State work, depth, synchronization, and buffering costs.

What This Part Should Cover

  • Power-of-two offsets
  • Reading a stable prior round
  • O(log n) parallel depth and O(n log n) work
  • Reverse traversal of the computation graph

Follow-up questions

  • How would you implement an exclusive scan?
  • When would a work-efficient tree scan be preferable?
  • How would floating-point associativity affect a parallel result?