OpenAI AI Engineer Interview Questions & Answers

import numpy as np


def masked_cross_entropy(
    logits: np.ndarray,
    labels: np.ndarray,
    mask: np.ndarray,
    smoothing: float = 0.0,
    temperature: float = 1.0,
    reduction: str = "mean",
) -> tuple[float | np.ndarray, np.ndarray]:
    """Compute masked label-smoothed cross-entropy and d(loss)/d(logits)."""

    # Validate that the three main inputs are NumPy arrays with the expected rank.
    if not isinstance(logits, np.ndarray) or logits.ndim != 2:
        raise ValueError("logits must be a 2D NumPy array")
    if not isinstance(labels, np.ndarray) or labels.ndim != 1:
        raise ValueError("labels must be a 1D NumPy array")
    if not isinstance(mask, np.ndarray) or mask.ndim != 1:
        raise ValueError("mask must be a 1D NumPy array")

    n, c = logits.shape

    # Cross-entropy needs at least one class column.
    if c <= 0:
        raise ValueError("logits must contain at least one class")

    # Every batch row must have exactly one label and one mask value.
    if labels.shape != (n,):
        raise ValueError("labels must have shape (N,)")
    if mask.shape != (n,):
        raise ValueError("mask must have shape (N,)")

    # Logits must use floating-point arithmetic for normalization and gradients.
    if not np.issubdtype(logits.dtype, np.floating):
        raise TypeError("logits must have a floating-point dtype")

    # Labels are class indices, so they must use an integer dtype.
    if not np.issubdtype(labels.dtype, np.integer):
        raise TypeError("labels must have an integer dtype")

    # The mask must be explicitly Boolean so row inclusion is unambiguous.
    if mask.dtype != np.bool_:
        raise TypeError("mask must have Boolean dtype")

    # NaN or infinity would make the probability calculation invalid.
    if not np.all(np.isfinite(logits)):
        raise ValueError("logits must contain only finite values")

    # Every label must identify a valid class column.
    if np.any(labels < 0) or np.any(labels >= c):
        raise ValueError("labels must be in [0, C-1]")

    # The diagram uses smoothing in [0, 1).
    if not np.isfinite(smoothing) or not 0.0 <= smoothing < 1.0:
        raise ValueError("smoothing must be in [0, 1)")

    # Temperature is a divisor, so it must be finite and strictly positive.
    if not np.isfinite(temperature) or temperature <= 0.0:
        raise ValueError("temperature must be finite and greater than 0")

    # Normalize the reduction name before checking the supported choices.
    if not isinstance(reduction, str):
        raise TypeError("reduction must be a string")
    reduction = reduction.lower()
    if reduction not in {"none", "sum", "mean"}:
        raise ValueError("reduction must be 'none', 'sum', or 'mean'")

    # Step 1: divide by temperature to control probability sharpness.
    z = logits / temperature

    # Step 2: subtract each row maximum before exponentiation for stability.
    row_max = np.max(z, axis=1, keepdims=True)
    shifted = z - row_max
    exp_shifted = np.exp(shifted)
    logsumexp = np.log(np.sum(exp_shifted, axis=1, keepdims=True))
    log_probs = shifted - logsumexp

    # Step 3: create one-hot targets with the same (N, C) shape as logits.
    one_hot = np.zeros_like(log_probs)
    one_hot[np.arange(n), labels] = 1.0

    # Step 4: mix the one-hot target with a uniform distribution over C classes.
    targets = (1.0 - smoothing) * one_hot + smoothing / c

    # Step 5: sum the class contributions to obtain one loss for each row.
    per_row_loss = -np.sum(targets * log_probs, axis=1)

    # Step 6: excluded rows contribute exactly zero to the loss.
    masked_loss = per_row_loss * mask
    valid_count = int(mask.sum())

    # Step 7: return the loss in the requested reduction form.
    if reduction == "none":
        loss: float | np.ndarray = masked_loss.copy()
    elif reduction == "sum":
        loss = float(masked_loss.sum())
    else:
        # An all-masked batch has zero numerator. Using 1 avoids division by zero.
        loss = float(masked_loss.sum() / max(valid_count, 1))

    # Stable log probabilities can be exponentiated to recover softmax p.
    probabilities = np.exp(log_probs)

    # Temperature scaling adds the 1 / temperature factor to dL/dlogits.
    gradient = (probabilities - targets) / temperature

    # Expand the row mask across classes so excluded rows have zero gradient.
    gradient *= mask[:, None]

    # Mean reduction averages only over the included rows.
    if reduction == "mean":
        gradient /= max(valid_count, 1)

    # Return the requested loss and an (N, C) gradient.
    return loss, gradient


def main() -> None:
    # Run the exact concrete example used by the approved diagram.
    logits = np.array(
        [[2.0, 0.0, -1.0], [0.5, 1.5, 0.0]],
        dtype=np.float64,
    )
    labels = np.array([0, 2], dtype=np.int64)
    mask = np.array([True, False], dtype=np.bool_)

    loss, gradient = masked_cross_entropy(
        logits=logits,
        labels=labels,
        mask=mask,
        smoothing=0.1,
        temperature=2.0,
        reduction="mean",
    )

    # Only row 0 contributes because row 1 is masked out.
    print(f"loss = {loss:.6f}")
    print("gradient =")
    print(gradient)


if __name__ == "__main__":
    main()

from __future__ import annotations

from dataclasses import dataclass
from typing import TypeAlias

Matrix: TypeAlias = list[list[float]]


@dataclass(frozen=True)
class ForwardCache:
    # Cache the tensors needed by backpropagation.
    # X is replicated, while each entry of w_local represents one device's W_p.
    x: Matrix
    w_local: list[Matrix]


def shape(matrix: Matrix) -> tuple[int, int]:
    # Return matrix dimensions and reject ragged input because matrix
    # multiplication requires every row to have the same number of columns.
    rows = len(matrix)
    cols = len(matrix[0]) if rows else 0
    if any(len(row) != cols for row in matrix):
        raise ValueError("Matrix rows must all have the same length")
    return rows, cols


def transpose(matrix: Matrix) -> Matrix:
    # Swap rows and columns. Backpropagation needs X^T and W_p^T.
    rows, cols = shape(matrix)
    return [[matrix[row][col] for row in range(rows)] for col in range(cols)]


def matmul(left: Matrix, right: Matrix) -> Matrix:
    # Perform ordinary matrix multiplication after checking the inner dimensions.
    left_rows, left_cols = shape(left)
    right_rows, right_cols = shape(right)
    if left_cols != right_rows:
        raise ValueError(
            f"Incompatible shapes: {(left_rows, left_cols)} and {(right_rows, right_cols)}"
        )

    # Each result cell is one row-by-column dot product.
    return [
        [sum(left[i][k] * right[k][j] for k in range(left_cols)) for j in range(right_cols)]
        for i in range(left_rows)
    ]


def column_shard_sizes(total_columns: int, devices: int) -> list[int]:
    # Use the diagram's uneven-shard rule. Every device first gets floor(N/P)
    # columns. The first N % P devices receive one additional column.
    if devices <= 0:
        raise ValueError("devices must be positive")
    if total_columns < 0:
        raise ValueError("total_columns cannot be negative")

    base = total_columns // devices
    remainder = total_columns % devices
    return [base + (1 if p < remainder else 0) for p in range(devices)]


def shard_columns(matrix: Matrix, sizes: list[int]) -> list[Matrix]:
    # Split a full matrix into contiguous column blocks using prefix boundaries.
    rows, cols = shape(matrix)
    if any(width < 0 for width in sizes):
        raise ValueError("Shard widths cannot be negative")
    if sum(sizes) != cols:
        raise ValueError("Shard sizes must sum to the matrix column count")

    shards: list[Matrix] = []
    start = 0
    for width in sizes:
        end = start + width
        shards.append([row[start:end] for row in matrix])
        start = end
    return shards


def concat_columns(shards: list[Matrix]) -> Matrix:
    # Concatenate Y_p blocks in device order. This is the local simulation
    # of the diagram's all-gather along the output-column dimension.
    if not shards:
        raise ValueError("At least one shard is required")

    row_count = shape(shards[0])[0]
    if any(shape(shard)[0] != row_count for shard in shards):
        raise ValueError("All column shards must have the same row count")

    result: Matrix = []
    for row_index in range(row_count):
        row: list[float] = []
        for shard in shards:
            row.extend(shard[row_index])
        result.append(row)
    return result


def all_reduce_sum(partials: list[Matrix], output_shape: tuple[int, int]) -> Matrix:
    # Sum overlapping dX_p tensors. This models an all-reduce with SUM.
    rows, cols = output_shape
    result: Matrix = [[0.0 for _ in range(cols)] for _ in range(rows)]

    for partial in partials:
        if shape(partial) != output_shape:
            raise ValueError("Every dX partial must have shape [B, K]")
        for i in range(rows):
            for j in range(cols):
                result[i][j] += partial[i][j]
    return result


def forward(x: Matrix, w_local: list[Matrix]) -> tuple[Matrix, ForwardCache]:
    # X is replicated on all devices. Each device computes only the output
    # columns owned by its local weight matrix W_p.
    if not w_local:
        raise ValueError("At least one weight shard is required")

    batch, input_features = shape(x)
    del batch  # The row count is carried by X and does not need a separate variable.

    for shard in w_local:
        shard_rows, _ = shape(shard)
        if shard_rows != input_features:
            raise ValueError("Every W_p must have shape [K, N_p]")

    y_local: list[Matrix] = [matmul(x, shard) for shard in w_local]

    # Concatenation in device order simulates the forward all-gather.
    y = concat_columns(y_local)
    return y, ForwardCache(x=x, w_local=w_local)


def backward(dy: Matrix, cache: ForwardCache) -> tuple[Matrix, list[Matrix]]:
    # Recover the replicated X and the original column-sharded weights.
    x = cache.x
    w_local = cache.w_local
    batch, input_features = shape(x)

    # Scatter dY by exactly the same column widths used for W.
    shard_sizes = [shape(shard)[1] for shard in w_local]
    if shape(dy) != (batch, sum(shard_sizes)):
        raise ValueError("dY must have global shape [B, N]")
    dy_local = shard_columns(dy, shard_sizes)

    x_t = transpose(x)
    dw_local: list[Matrix] = []
    dx_local: list[Matrix] = []

    for w_shard, dy_shard in zip(w_local, dy_local, strict=True):
        _, local_columns = shape(w_shard)

        # dW_p = X^T @ dY_p. It has shape [K, N_p] and stays with W_p.
        dw_local.append(matmul(x_t, dy_shard))

        if local_columns == 0:
            # When N < P, a device may own no columns. It then contributes
            # zero to dX and owns an empty [K, 0] dW shard.
            dx_local.append([[0.0 for _ in range(input_features)] for _ in range(batch)])
        else:
            # dX_p = dY_p @ W_p^T. Every non-empty shard contributes to
            # the same global [B, K] input-gradient tensor.
            dx_local.append(matmul(dy_shard, transpose(w_shard)))

    # Sum every partial dX_p. This simulates all-reduce(SUM).
    dx = all_reduce_sum(dx_local, (batch, input_features))
    return dx, dw_local


def matrices_close(
    left: Matrix,
    right: Matrix,
    tolerance: float = 1e-9,
) -> bool:
    # Compare sharded and reference calculations with floating-point tolerance.
    if shape(left) != shape(right):
        return False

    rows, cols = shape(left)
    return all(abs(left[i][j] - right[i][j]) <= tolerance for i in range(rows) for j in range(cols))


def run_reference_test(
    x: Matrix,
    w: Matrix,
    dy: Matrix,
    devices: int,
) -> None:
    # Build contiguous W_p shards using the same boundaries that will also
    # be used to scatter dY during backpropagation.
    _, total_columns = shape(w)
    sizes = column_shard_sizes(total_columns, devices)
    w_local = shard_columns(w, sizes)

    # Run the sharded forward path and assemble the full Y.
    y, cache = forward(x, w_local)

    # Run the sharded backward path and all-reduce the partial dX tensors.
    dx, dw_local = backward(dy, cache)

    # Compute the direct unsharded reference requested by the diagram.
    y_ref = matmul(x, w)
    dw_ref = matmul(transpose(x), dy)
    dx_ref = matmul(dy, transpose(w))
    dw_ref_local = shard_columns(dw_ref, sizes)

    # The distributed decomposition is correct only if all results match
    # the corresponding unsharded computation.
    assert matrices_close(y, y_ref)
    assert matrices_close(dx, dx_ref)
    assert len(dw_local) == len(dw_ref_local)
    for actual, expected in zip(dw_local, dw_ref_local, strict=True):
        assert matrices_close(actual, expected)


def main() -> None:
    # The diagram supplies symbolic shapes rather than fixed numeric values.
    # These deterministic fixtures test the same forward/backward equations.
    x: Matrix = [
        [1.0, 2.0, -1.0],
        [0.5, -2.0, 3.0],
    ]
    w: Matrix = [
        [1.0, 0.0, 2.0, -1.0, 3.0],
        [2.0, 1.0, 0.0, 4.0, -2.0],
        [-1.0, 3.0, 1.0, 0.5, 2.0],
    ]
    dy: Matrix = [
        [1.0, -1.0, 0.5, 2.0, 0.0],
        [0.0, 2.0, -1.0, 1.5, 3.0],
    ]

    # N=5 and P=2 exercises the diagram's uneven-shard rule: [3, 2].
    run_reference_test(x, w, dy, devices=2)

    # P>N exercises the diagram's N<P edge case. Some devices own zero columns.
    run_reference_test(x, w, dy, devices=7)

    print("All sharded results match the unsharded reference.")


if __name__ == "__main__":
    main()