Implement a numerically stable backward pass for layer normalization from scratch in NumPy.

Constraints:
- Input: x of shape (B, T, D)
- Parameters: gamma, beta of shape (D,)
- Forward:
    y = gamma * (x - mean) / sqrt(var + eps) + beta

Requirements:
1. Derive and implement gradients w.r.t. x, gamma, beta manually (no autodiff).
2. Avoid redundant recomputation — reuse intermediates where possible.
3. Ensure numerical stability (discuss where instability can occur).
4. Provide a gradient check using finite differences.
5. Analyze time and memory complexity.
6. Explain how you would fuse this into a single kernel for GPU execution.

Do not use PyTorch, TensorFlow, JAX, or autograd.


The user wants me to implement a numerically stable backward pass for layer normalization from scratch in NumPy. Let me break down what I need to do:

1. **Forward pass**: y = gamma * (x - mean) / sqrt(var + eps) + beta
   - Input shape: (B, T, D) where B=batch, T=sequence length, D=hidden dimension
   - mean and var are computed over the last dimension (D)

2. **Backward pass**: I need to derive gradients for:
   - dx (gradient w.r.t. input x)
   - d_gamma (gradient w.r.t. gamma)
   - d_beta (gradient w.r.t. beta)

Let me derive the gradients:

**Forward pass:**
- mean = (1/D) * sum(x, axis=-1), shape (B, T,)
- var = (1/D) * sum((x - mean)^2, axis=-1), shape (B, T,)
- x_centered = x - mean (broadcasts)
- x_norm = x_centered / sqrt(var + eps)
- y = gamma * x_norm + beta

**Backward pass:**

Let dy be the incoming gradient (same shape as y).

1. **d_beta**: This is straightforward
   d_beta = sum(dy, axis=-1) → shape (D,)

2. **d_gamma**:
   dy/d_gamma = x_norm
   d_gamma = sum(dy * x_norm, axis=-1) → shape (D,)

3. **dx**: This is the tricky part
   We have: y = gamma * (x - mean) / sqrt(var + eps) + beta

   Let std = sqrt(var + eps)
   x_norm = (x - mean) / std

   dy/dx = gamma / std

   But there's a coupling through mean and var...

   Breaking this down: y_i = gamma_i * (x_i - mean) / std + beta_i where i indexes the D dimension. The loss gradient chains through: dl/dx_i = gamma_i / std * dl/dy_i, plus additional terms from how mean and var depend on x. Computing d_mean/dx_j = 1/D and d_var/dx_j = 2/D * (x_j - mean), I can express the full derivative as dl/dx_j = gamma_j / std * dl/dy_j - (gamma_j / (D * std)) * sum(dl/dy_i) - (gamma_j * (x_j - mean) / (D * std^3)) * sum(dl/dy_i * (x_i - mean)).

Let me define the normalized input z_i = (x_i - mean) / std so y_i = gamma_i * z_i + beta_i. Then dl/dz_i = gamma_i * dl/dy_i, and dl/dx_j = (1/std) * dl/dz_j - (1/(D * std)) * sum_i(dl/dz_i) - (z_j / (D * std)) * sum_i(dl/dz_i * z_i). This simplifies to the formula above when I substitute back the definitions of z and std.

For efficient implementation, I need to compute z, mean, and var during the forward pass, then reuse these values when computing gradients. The backward pass requires accumulating dl/dz across the batch and sequence dimensions for each feature, computing the weighted sum of gradients scaled by z, and finally combining these terms to get dl/dx. 1) * sum_dy_centered / D
- For instability: eps should be large enough to prevent sqrt(0), but std is already computed with eps in forward pass, so this is handled. The main risk is in computing std^3 in the backward pass, but since std ≥ sqrt(eps) > 0, this remains numerically stable. I should test with various epsilon values to verify.

For the gradient check, I'll use finite differences with a small perturbation delta to verify the computed gradients against numerical approximations.

The computational complexity breaks down to O(BTD) for the forward pass, while the backward pass requires an additional O(BTD) for dx, dgamma, and dbeta calculations. Memory-wise, I need to store x, mean, var, and normalized values, each taking O(BTD) space, plus O(D) for the parameters themselves. For GPU optimization, I'd restructure the kernel to compute statistics and normalization in a single fused operation with minimal global memory access.

The gradient computation follows a similar pattern but needs careful handling of reduction operations—keeping everything in registers or shared memory, precomputing reduction sums, and managing partial reductions across threads before writing results back to global memory.

Looking at the layer normalization problem, I need to implement both forward and backward passes with proper gradient derivation, numerical stability, and verification.

Let me start by implementing the complete solution:

write

"""
Layer Normalization from scratch in NumPy
==========================================
Numerically stable backward pass implementation with gradient checking.

Author: pi coding assistant
Date: 2026-04-22
"""

import numpy as np
from typing import Tuple, Dict, Optional
import time

# =============================================================================
# Numerical constants
# =============================================================================
DEFAULT_EPS = 1e-8

# =============================================================================
# Helper functions
# =============================================================================

def logsumexp(x: np.ndarray, axis: int = -1, keepdims: bool = True) -> np.ndarray:
    """Numerically stable log-sum-exp."""
    max_x = np.max(x, axis=axis, keepdims=True)
    return max_x + np.log(np.exp(x - max_x).sum(axis=axis, keepdims=True))

# =============================================================================
# Layer Normalization Forward Pass
# =============================================================================

def layer_norm_forward(
    x: np.ndarray,
    gamma: np.ndarray,
    beta: np.ndarray,
    eps: float = DEFAULT_EPS
) -> Tuple[np.ndarray, Dict]:
    """
    Layer Norm forward pass.

    y = gamma * (x - mean) / sqrt(var + eps) + beta

    Args:
        x: Input tensor of shape (B, T, D)
        gamma: Scale parameter of shape (D,)
        beta: Bias parameter of shape (D,)
        eps: Small constant for numerical stability

    Returns:
        y: Normalized output of shape (B, T, D)
        cache: Dictionary of intermediates for backward pass
    """
    B, T, D = x.shape

    # Compute mean over feature dimension
    # mean[b, t] = (1/D) * sum_d x[b, t, d]
    mean = np.mean(x, axis=-1, keepdims=True)  # (B, T, 1)

    # Compute variance over feature dimension
    # var[b, t] = (1/D) * sum_d (x[b, t, d] - mean[b, t])^2
    x_centered = x - mean  # (B, T, D)
    var = np.mean(x_centered ** 2, axis=-1, keepdims=True)  # (B, T, 1)

    # Compute standard deviation with eps for numerical stability
    # std >= sqrt(eps) > 0, preventing division by zero
    std = np.sqrt(var + eps)  # (B, T, 1)

    # Normalize
    x_norm = x_centered / std  # (B, T, D)

    # Scale and shift
    y = gamma * x_norm + beta  # (B, T, D)

    # Cache intermediates for backward pass
    # We store only what we need to avoid recomputation
    cache = {
        'x': x,              # Original input (needed for gradient check)
        'x_centered': x_centered,
        'x_norm': x_norm,    # Normalized values (needed for d_gamma)
        'mean': mean,        # (B, T, 1)
        'var': var,          # (B, T, 1)
        'std': std,          # (B, T, 1)
        'gamma': gamma,      # Needed for gradient computation
        'beta': beta,        # Needed for gradient check
        'eps': eps,
        'B': B,
        'T': T,
        'D': D
    }

    return y, cache

# =============================================================================
# Layer Normalization Backward Pass
# =============================================================================

def layer_norm_backward(
    dy: np.ndarray,
    cache: Dict
) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
    """
    Layer Norm backward pass.

    Derivation:
    -----------
    Let:
        μ = mean[x]  (over axis=-1)
        σ² = var[x]  (over axis=-1)
        σ = sqrt(σ² + eps)
        x̄ = (x - μ) / σ  (normalized)
        y = γ * x̄ + β

    Then:
        ∂L/∂γ = sum(dy * x̄) over (B, T)
        ∂L/∂β = sum(dy) over (B, T)

    For ∂L/∂x:
        ∂L/∂x_i = (∂L/∂x̄_i / σ)
                 - (∂L/∂x̄_i / D) * (1/σ)
                 - (x̄_i / D) * (∂L/∂σ²) * (2/σ)

        where ∂L/∂σ² = -0.5 * sum_i(∂L/∂x̄_i * x̄_i) / σ³

    Consolidating:
        ∂L/∂x_i = (γ_i / σ) * [∂L/∂y_i
                               - mean(∂L/∂y)
                               - x̄_i * mean(∂L/∂y * x̄)]

    Args:
        dy: Upstream gradient of shape (B, T, D)
        cache: Forward pass intermediates

    Returns:
        dx: Gradient w.r.t. input x, shape (B, T, D)
        d_gamma: Gradient w.r.t. gamma, shape (D,)
        d_beta: Gradient w.r.t. beta, shape (D,)
    """
    x = cache['x']
    x_centered = cache['x_centered']
    x_norm = cache['x_norm']
    mean = cache['mean']
    std = cache['std']
    gamma = cache['gamma']
    eps = cache['eps']
    B, T, D = cache['B'], cache['T'], cache['D']

    # -------------------------------------------------------------------------
    # 1. Compute gradients w.r.t. gamma and beta
    # -------------------------------------------------------------------------
    # d_gamma[d] = sum_{b,t} dy[b,t,d] * x_norm[b,t,d]
    d_gamma = np.sum(dy * x_norm, axis=(0, 1))  # (D,)

    # d_beta[d] = sum_{b,t} dy[b,t,d]
    d_beta = np.sum(dy, axis=(0, 1))  # (D,)

    # -------------------------------------------------------------------------
    # 2. Compute gradient w.r.t. normalized input
    # -------------------------------------------------------------------------
    # dz = dy * gamma  (chain rule: y = gamma * x_norm + beta)
    # Note: We can compute this and reuse in dx computation
    dz = dy * gamma  # (B, T, D)

    # -------------------------------------------------------------------------
    # 3. Compute gradient w.r.t. x
    # -------------------------------------------------------------------------
    #
    # From the derivation:
    # dx = (dz - mean(dz) - x_norm * mean(dz * x_norm)) / std
    #
    # This comes from applying the chain rule considering:
    # - Direct dependence of x on x_norm through (x - mean) / std
    # - Indirect dependence through mean and std
    #
    # Key insight: We compute the two reduction terms efficiently:
    # - mean(dz) = (1/D) * sum(dz, axis=-1, keepdims=True)
    # - mean(dz * x_norm) = (1/D) * sum(dz * x_norm, axis=-1, keepdims=True)
    #

    # Compute reduction terms (these are O(BTD) each)
    sum_dz = np.sum(dz, axis=-1, keepdims=True)  # (B, T, 1)
    sum_dz_xnorm = np.sum(dz * x_norm, axis=-1, keepdims=True)  # (B, T, 1)

    # Compute dx using the consolidated formula
    # dx = (dz - (sum_dz / D) - x_norm * (sum_dz_xnorm / D)) / std
    dx = (dz - sum_dz / D - x_norm * sum_dz_xnorm / D) / std

    # -------------------------------------------------------------------------
    # Numerical stability analysis:
    # ---------------------------
    # 1. Division by std: We use std = sqrt(var + eps), so std >= sqrt(eps) > 0
    #    Example: eps = 1e-8 => std >= 1e-4, so no division by zero
    #
    # 2. Division by D: D is typically 512-4096, so this is stable
    #
    # 3. The formula (dz - mean(dz) - x_norm * mean(dz * x_norm)) / std
    #    When std is very small, the gradient can be large, but this is
    #    mathematically correct - small std means large normalization effect
    #
    # 4. For extreme stability, we could use the two-pass formula:
    #    dx = (dz - mean(dz) - x_norm * mean(dz * x_norm))
    #    dx = dx / std
    #    This avoids any intermediate overflow/underflow in the subtraction
    #
    # 5. Alternative numerically stable computation using centering trick:
    #    temp = dz / std
    #    dx = temp - x_norm * (sum(temp * x_norm) / D)
    #    dx = dx - sum(dx) / D  (but this is less efficient)
    #
    # 6. Catastrophic cancellation can occur in: (dz - mean(dz))
    #    When dz is roughly constant across D, mean(dz) ≈ dz, causing
    #    cancellation. However, this is exactly when dx should be small,
    #    so the cancellation is benign (relative error is small).
    #
    # 7. The x_norm * mean(dz * x_norm) term can also suffer from cancellation
    #    when mean(dz * x_norm) ≈ 0, but again this is when the term is small.
    #
    # 8. For fp16 or extreme cases, consider pairwise summation for reductions
    #    and/or higher precision accumulators.
    #
    # -------------------------------------------------------------------------

    return dx, d_gamma, d_beta

# =============================================================================
# Layer Norm Module (combines forward and backward)
# =============================================================================

class LayerNorm:
    """
    Layer Normalization module with manual gradient computation.

    Forward: y = gamma * (x - mean) / sqrt(var + eps) + beta
    Backward: Computes gradients w.r.t. x, gamma, beta
    """

    def __init__(self, normalized_shape: int, eps: float = DEFAULT_EPS):
        """
        Args:
            normalized_shape: Dimension D of the feature space
            eps: Epsilon for numerical stability in sqrt(var + eps)
        """
        self.normalized_shape = normalized_shape
        self.eps = eps

        # Initialize gamma (scale) and beta (shift) parameters
        # Xavier initialization for gamma to keep variance stable
        self.gamma = np.ones(normalized_shape)  # Scale initialized to 1
        self.beta = np.zeros(normalized_shape)   # Shift initialized to 0

        # Storage for gradients
        self.d_gamma = None
        self.d_beta = None

    def forward(self, x: np.ndarray) -> Tuple[np.ndarray, Dict]:
        """Forward pass."""
        assert x.shape[-1] == self.normalized_shape, \
            f"Expected last dimension {self.normalized_shape}, got {x.shape[-1]}"
        return layer_norm_forward(x, self.gamma, self.beta, self.eps)

    def backward(self, dy: np.ndarray, cache: Dict) -> np.ndarray:
        """
        Backward pass.

        Args:
            dy: Upstream gradient of shape (B, T, D)
            cache: Forward pass cache

        Returns:
            dx: Gradient w.r.t. input x
        """
        dx, self.d_gamma, self.d_beta = layer_norm_backward(dy, cache)
        return dx

    def parameters(self) -> Tuple[np.ndarray, np.ndarray]:
        """Return (gamma, beta)."""
        return self.gamma, self.beta

    def gradients(self) -> Tuple[np.ndarray, np.ndarray]:
        """Return (d_gamma, d_beta)."""
        return self.d_gamma, self.d_beta

# =============================================================================
# Gradient Checking via Finite Differences
# =============================================================================

def compute_numerical_gradient(
    func: callable,
    x: np.ndarray,
    param: np.ndarray,
    param_name: str,
    dy: np.ndarray,
    h: float = 1e-5,
    cache_base: Optional[Dict] = None
) -> np.ndarray:
    """
    Compute numerical gradient using finite differences.

    Uses central difference: df/dx ≈ (f(x+h) - f(x-h)) / (2h)

    Args:
        func: Function that takes (x, param) and returns (y, cache)
        x: Input tensor
        param: Parameter array to compute gradient for
        param_name: Name of parameter ('gamma' or 'beta')
        dy: Upstream gradient
        h: Step size for finite differences
        cache_base: Base cache if modifying param

    Returns:
        Numerical gradient of same shape as param
    """
    num_grad = np.zeros_like(param)

    # Flatten for easier iteration
    param_flat = param.flatten()
    num_grad_flat = num_grad.flatten()

    for i in range(len(param_flat)):
        # Save original value
        original = param_flat[i]

        # f(x, param + h)
        param_flat[i] = original + h
        y_plus, _ = func(x, param)
        loss_plus = np.sum(y_plus * dy)

        # f(x, param - h)
        param_flat[i] = original - h
        y_minus, _ = func(x, param)
        loss_minus = np.sum(y_minus * dy)

        # Central difference
        num_grad_flat[i] = (loss_plus - loss_minus) / (2 * h)

        # Restore original
        param_flat[i] = original

    return num_grad_flat.reshape(param.shape)

def compute_numerical_gradient_wrt_x(
    model: LayerNorm,
    x: np.ndarray,
    dy: np.ndarray,
    h: float = 1e-5
) -> np.ndarray:
    """
    Compute numerical gradient w.r.t. input x.

    For x, we perturb each element and compute the effect on the loss.
    Since x is large (B*T*D), we compute gradient at element level carefully.
    """
    B, T, D = x.shape
    num_grad = np.zeros_like(x)

    # Flatten x for easier iteration (but be careful with large tensors)
    x_flat = x.flatten()
    num_grad_flat = num_grad.flatten()

    total_elements = len(x_flat)
    # For very large tensors, we might want to subsample
    # Here we do full computation for correctness

    print(f"  Computing numerical gradient for x ({total_elements} elements)...")

    for i in range(total_elements):
        if (i + 1) % 10000 == 0:
            print(f"    Progress: {i+1}/{total_elements}")

        original = x_flat[i]

        # f(x + h)
        x_flat[i] = original + h
        y_plus, _ = model.forward(x)
        loss_plus = np.sum(y_plus * dy)

        # f(x - h)
        x_flat[i] = original - h
        y_minus, _ = model.forward(x)
        loss_minus = np.sum(y_minus * dy)

        # Central difference
        num_grad_flat[i] = (loss_plus - loss_minus) / (2 * h)

        # Restore
        x_flat[i] = original

    return num_grad_flat.reshape(x.shape)

def gradient_check(
    model: LayerNorm,
    x: np.ndarray,
    dy: np.ndarray,
    rtol: float = 1e-4,
    atol: float = 1e-5,
    verbose: bool = True
) -> Dict[str, bool]:
    """
    Perform gradient check for all parameters.

    Checks: |analytical - numerical| <= atol + rtol * |numerical|
    """
    results = {}

    # Store original gamma and beta
    orig_gamma = model.gamma.copy()
    orig_beta = model.beta.copy()

    # Helper function for parameter gradient checking
    def func_with_param(x_input, param):
        # Temporarily set the parameter
        if param is model.gamma:
            model.gamma = param
        else:
            model.beta = param
        y, cache = model.forward(x_input)
        # Restore original
        if param is model.gamma:
            model.gamma = orig_gamma
        else:
            model.beta = orig_beta
        return y, cache

    # -------------------------------------------------------------------------
    # Check gradient w.r.t. gamma
    # -------------------------------------------------------------------------
    if verbose:
        print("\n" + "="*60)
        print("GRADIENT CHECK: gamma")
        print("="*60)

    model.gamma = orig_gamma.copy()
    y, cache = model.forward(x)
    dx, d_gamma_analytical, _ = model.backward(dy, cache)

    d_gamma_numerical = compute_numerical_gradient(
        func_with_param, x, orig_gamma.copy(), 'gamma', dy
    )

    # Reset gamma
    model.gamma = orig_gamma.copy()

    # Compute relative error
    diff = np.abs(d_gamma_analytical - d_gamma_numerical)
    tolerance = atol + rtol * np.abs(d_gamma_numerical)
    passed = np.all(diff <= tolerance)

    if verbose:
        print(f"Analytical gradient shape: {d_gamma_analytical.shape}")
        print(f"Numerical gradient shape: {d_gamma_numerical.shape}")
        print(f"Max absolute difference: {np.max(diff):.2e}")
        print(f"Max relative tolerance: {np.max(tolerance):.2e}")
        print(f"Mean analytical: {np.mean(np.abs(d_gamma_analytical)):.6e}")
        print(f"Mean numerical: {np.mean(np.abs(d_gamma_numerical)):.6e}")
        print(f"\nGradient check: {'PASSED ✓' if passed else 'FAILED ✗'}")

    results['gamma'] = passed

    # -------------------------------------------------------------------------
    # Check gradient w.r.t. beta
    # -------------------------------------------------------------------------
    if verbose:
        print("\n" + "="*60)
        print("GRADIENT CHECK: beta")
        print("="*60)

    model.beta = orig_beta.copy()
    y, cache = model.forward(x)
    dx, _, d_beta_analytical = model.backward(dy, cache)

    d_beta_numerical = compute_numerical_gradient(
        func_with_param, x, orig_beta.copy(), 'beta', dy
    )

    # Reset beta
    model.beta = orig_beta.copy()

    diff = np.abs(d_beta_analytical - d_beta_numerical)
    tolerance = atol + rtol * np.abs(d_beta_numerical)
    passed = np.all(diff <= tolerance)

    if verbose:
        print(f"Analytical gradient shape: {d_beta_analytical.shape}")
        print(f"Numerical gradient shape: {d_beta_numerical.shape}")
        print(f"Max absolute difference: {np.max(diff):.2e}")
        print(f"Max relative tolerance: {np.max(tolerance):.2e}")
        print(f"Mean analytical: {np.mean(np.abs(d_beta_analytical)):.6e}")
        print(f"Mean numerical: {np.mean(np.abs(d_beta_numerical)):.6e}")
        print(f"\nGradient check: {'PASSED ✓' if passed else 'FAILED ✗'}")

    results['beta'] = passed

    # -------------------------------------------------------------------------
    # Check gradient w.r.t. x
    # -------------------------------------------------------------------------
    if verbose:
        print("\n" + "="*60)
        print("GRADIENT CHECK: x (input)")
        print("="*60)

    # Need to recompute since model was modified
    model.gamma = orig_gamma.copy()
    model.beta = orig_beta.copy()
    y, cache = model.forward(x)
    dx_analytical, _, _ = model.backward(dy, cache)

    # Numerical gradient for x is expensive - only do for smaller tensors
    total_elements = x.size
    if total_elements <= 100000:  # Only check if reasonably sized
        d_x_numerical = compute_numerical_gradient_wrt_x(model, x, dy)

        diff = np.abs(dx_analytical - d_x_numerical)
        tolerance = atol + rtol * np.abs(d_x_numerical)
        passed = np.all(diff <= tolerance)

        if verbose:
            print(f"Analytical gradient shape: {dx_analytical.shape}")
            print(f"Numerical gradient shape: {d_x_numerical.shape}")
            print(f"Max absolute difference: {np.max(diff):.2e}")
            print(f"Max relative tolerance: {np.max(tolerance):.2e}")
            print(f"Mean analytical: {np.mean(np.abs(dx_analytical)):.6e}")
            print(f"Mean numerical: {np.mean(np.abs(d_x_numerical)):.6e}")
            print(f"\nGradient check: {'PASSED ✓' if passed else 'FAILED ✗'}")

        results['x'] = passed
    else:
        if verbose:
            print(f"Skipping x gradient check (tensor too large: {total_elements} elements)")
            print("Performing spot check on subset...")

        # Spot check on subset
        indices = [tuple(np.random.randint(s) for s in x.shape) for _ in range(100)]
        model.gamma = orig_gamma.copy()
        model.beta = orig_beta.copy()
        y, cache = model.forward(x)
        dx_analytical_spot, _, _ = model.backward(dy, cache)

        d_x_numerical_spot = np.zeros((100,))
        dx_analytical_spot_vals = np.zeros((100,))

        for idx, (bi, ti, di) in enumerate(indices):
            orig_val = x[bi, ti, di]

            # x + h
            x[bi, ti, di] = orig_val + 1e-5
            y_plus, _ = model.forward(x)
            loss_plus = np.sum(y_plus * dy)

            # x - h
            x[bi, ti, di] = orig_val - 1e-5
            y_minus, _ = model.forward(x)
            loss_minus = np.sum(y_minus * dy)

            d_x_numerical_spot[idx] = (loss_plus - loss_minus) / 2e-5
            dx_analytical_spot_vals[idx] = dx_analytical_spot[bi, ti, di]

            # Restore
            x[bi, ti, di] = orig_val

        diff = np.abs(dx_analytical_spot_vals - d_x_numerical_spot)
        tolerance = atol + rtol * np.abs(d_x_numerical_spot)
        passed = np.all(diff <= tolerance)

        if verbose:
            print(f"Spot check (100 random elements)")
            print(f"Max absolute difference: {np.max(diff):.2e}")
            print(f"Max relative tolerance: {np.max(tolerance):.2e}")
            print(f"\nSpot check: {'PASSED ✓' if passed else 'FAILED ✗'}")

        results['x'] = passed

    # Restore original parameters
    model.gamma = orig_gamma.copy()
    model.beta = orig_beta.copy()

    return results

# =============================================================================
# Complexity Analysis
# =============================================================================

def analyze_complexity():
    """Print complexity analysis for layer norm forward and backward."""
    print("""
╔══════════════════════════════════════════════════════════════════════════════╗
║                    COMPLEXITY ANALYSIS                                         ║
╠══════════════════════════════════════════════════════════════════════════════╣
║                                                                              ║
║  Input: x ∈ ℝ^(B×T×D)                                                        ║
║  Parameters: γ, β ∈ ℝ^D                                                      ║
║                                                                              ║
╠══════════════════════════════════════════════════════════════════════════════╣
║  FORWARD PASS                                                                ║
╠══════════════════════════════════════════════════════════════════════════════╣
║                                                                              ║
║  Operation              │ Work (FLOPs)      │ Memory                        ║
║  ────────────────────────────────────────────────────────────────────────── ║
║  mean (reduction)       │ O(BTD)            │ -                             ║
║  x - mean (broadcast)   │ O(BTD)            │ O(BTD)                        ║
║  var (reduction)        │ O(BTD)            │ -                             ║
║  sqrt(var + eps)        │ O(BT)             │ -                             ║
║  divide by std          │ O(BTD)            │ -                             ║
║  gamma * x_norm         │ O(BTD)            │ -                             ║
║  add beta               │ O(BTD)            │ O(BTD) output                 ║
║  ────────────────────────────────────────────────────────────────────────── ║
║  TOTAL                    │ 5×O(BTD)          │ O(BTD)                       ║
║                                                                              ║
╠══════════════════════════════════════════════════════════════════════════════╣
║  BACKWARD PASS                                                              ║
╠══════════════════════════════════════════════════════════════════════════════╣
║                                                                              ║
║  Operation              │ Work (FLOPs)      │ Memory                        ║
║  ────────────────────────────────────────────────────────────────────────── ║
║  d_gamma = sum(dy*x_norm)│ O(BTD)           │ O(D)                          ║
║  d_beta = sum(dy)       │ O(BTD)            │ O(D)                          ║
║  dz = dy * gamma        │ O(BTD)            │ O(BTD) (can be avoided)        ║
║  sum(dz)                │ O(BTD)            │ -                             ║
║  sum(dz * x_norm)       │ O(BTD)            │ -                             ║
║  dx computation         │ O(BTD)            │ O(BTD) output                 ║
║  ────────────────────────────────────────────────────────────────────────── ║
║  TOTAL                    │ 5×O(BTD)          │ O(BTD) + O(D)                ║
║                                                                              ║
╠══════════════════════════════════════════════════════════════════════════════╣
║  SUMMARY                                                                    ║
╠══════════════════════════════════════════════════════════════════════════════╣
║                                                                              ║
║  Time Complexity:  O(BTD) for both forward and backward                      ║
║  Space Complexity: O(BTD) for storing activations (during training)          ║
║                   O(BTD) during inference (no need to store)                 ║
║                                                                              ║
║  Cache efficiency: We store x_centered, x_norm, mean, std                    ║
║                   These are O(BTD) total, reused across all gradient comps   ║
║                                                                              ║
╚══════════════════════════════════════════════════════════════════════════════╝
""")

# =============================================================================
# GPU Kernel Fusion Design
# =============================================================================

def explain_gpu_fusion():
    """Explain GPU kernel fusion for layer normalization."""
    print("""
╔══════════════════════════════════════════════════════════════════════════════╗
║               GPU KERNEL FUSION DESIGN                                       ║
╠══════════════════════════════════════════════════════════════════════════════╣
║                                                                              ║
║  CURRENT SEPARATE KERNEL APPROACH:                                           ║
║  ────────────────────────────────────                                        ║
║  Kernel 1: Compute mean (reduction over D)                                   ║
║  Kernel 2: Compute variance (reduction over D)                               ║
║  Kernel 3: Normalize (element-wise)                                          ║
║  Kernel 4: Scale and shift (element-wise)                                    ║
║  Kernel 5: Backward kernels (x, gamma, beta)                                 ║
║                                                                              ║
║  Issues with separate kernels:                                               ║
║  • Multiple kernel launches (overhead)                                       ║
║  • Data movement between global memory passes                                 ║
║  • Can't use persistent threads for reduction efficiency                      ║
║                                                                              ║
╠══════════════════════════════════════════════════════════════════════════════╣
║  FUSED KERNEL DESIGN (Forward):                                              ║
╠══════════════════════════════════════════════════════════════════════════════╣
║                                                                              ║
║  Grid: (B × T) blocks, each handling one (b,t) position                      ║
║  Block: 256-512 threads                                                      ║
║                                                                              ║
║  PHASE 1: Load and compute local sum                                         ║
║  ─────────────────────────────────                                           ║
║  • Each thread loads x[b,t,d] into shared memory                             ║
║  • Compute partial sum using warp-level reduction                            ║
║  • Single thread writes mean to __shared__                                    ║
║                                                                              ║
║  PHASE 2: Compute variance locally                                           ║
║  ─────────────────────────────────                                           ║
║  • Re-load x with loaded mean                                                ║
║  • Compute (x-mean)² and partial variance                                    ║
║  • Reduce to get variance                                                   ║
║                                                                              ║
║  PHASE 3: Normalize and output                                               ║
║  ─────────────────────────────────                                           ║
║  • All threads compute: y = gamma * (x-mean) / sqrt(var+eps) + beta          ║
║  • Write to output (fully coalesced)                                         ║
║                                                                              ║
║  MEMORY ACCESS PATTERN:                                                      ║
║  • Each block reads contiguous D elements (coalesced)                        ║
║  • Use shared memory for intermediate results                                ║
║  • Output writes are also coalesced                                          ║
║                                                                              ║
╠══════════════════════════════════════════════════════════════════════════════╣
║  FUSED KERNEL DESIGN (Backward):                                             ║
╠══════════════════════════════════════════════════════════════════════════════╣
║                                                                              ║
║  Key insight: dz = dy * gamma can be merged into the computation              ║
║                                                                              ║
║  Grid: (B × T) blocks                                                       ║
║  Block: 256-512 threads                                                      ║
║                                                                              ║
║  SHARED MEMORY STRUCTURE:                                                    ║
║  [x_norm_0, x_norm_1, ..., x_norm_{D-1}]                                     ║
║  [dz_0, dz_1, ..., dz_{D-1}]                                                 ║
║                                                                              ║
║  ALGORITHM:                                                                  ║
║  1. Load x_norm and compute local dz = dy * gamma                            ║
║  2. Reduce to get sum(dz) and sum(dz * x_norm)                               ║
║  3. Second pass to compute dx using the formula:                             ║
║     dx = (dz - mean(dz) - x_norm * mean(dz*x_norm)) / std                    ║
║                                                                              ║
║  REDUCTION OPTIMIZATIONS:                                                    ║
║  • Warp-level shuffle reductions (no shared memory needed)                    ║
║  • Block-level using shared memory with tree reduction                        ║
║  • Use block-level primitives for final reduction                            ║
║                                                                              ║
╠══════════════════════════════════════════════════════════════════════════════╣
║  BENEFITS OF FUSION:                                                          ║
╠══════════════════════════════════════════════════════════════════════════════╣
║                                                                              ║
║  ✓ Reduced kernel launch overhead (1 kernel vs 4-5)                           ║
║  ✓ Better memory bandwidth utilization (single read, single write)          ║
║  ✓ Improved cache locality (data stays in registers/shared mem)              ║
║  ✓ Only loads x once, computes mean and var from same data                   ║
║  ✓ Backward can reuse cached values from forward (if memory allows)          ║
║  ✓ Lower register pressure allows for larger block sizes                     ║
║                                                                              ║
╠══════════════════════════════════════════════════════════════════════════════╣
║  CUDA KERNEL SKETCH (Pseudo-code):                                            ║
╠══════════════════════════════════════════════════════════════════════════════╣
║                                                                              ║
║  __global__ void layer_norm_fwd(float* y, const float* x,                    ║
║                                  const float* gamma, const float* beta,      ║
║                                  int B, int T, int D, float eps) {           ║
║                                                                              ║
║      __shared__ float mean_smem[256]; // block-level mean                     ║
║      __shared__ float var_smem[256];  // block-level variance                 ║
║      __shared__ float std_smem[256]; // block-level std                      ║
║                                                                              ║
║      int tid = threadIdx.x;                                                  ║
║      int bid = blockIdx.x;                                                   ║
║      int D_blk = (D + blockDim.x - 1) / blockDim.x;                          ║
║                                                                              ║
║      // Phase 1: Load and compute mean                                        ║
║      float sum = 0.0;                                                        ║
║      for (int i = 0; i < D_blk; i++) {                                       ║
║          int idx = bid * D + i * blockDim.x + tid;                           ║
║          sum += x[idx];                                                      ║
║      }                                                                       ║
║      sum = warpReduceSum(sum);                                               ║
║      if (tid % 32 == 0) mean_smem[tid / 32] = sum;                           ║
║      __syncthreads();                                                        ║
║                                                                              ║
║      float mean = mean_smem[0] / D;                                         ║
║                                                                              ║
║      // Phase 2: Compute variance                                            ║
║      sum = 0.0;                                                             ║
║      for (int i = 0; i < D_blk; i++) {                                       ║
║          int idx = bid * D + i * blockDim.x + tid;                           ║
║          float diff = x[idx] - mean;                                         ║
║          sum += diff * diff;                                                 ║
║      }                                                                       ║
║      sum = warpReduceSum(sum);                                               ║
║      if (tid % 32 == 0) var_smem[tid / 32] = sum;                           ║
║      __syncthreads();                                                       ║
║                                                                              ║
║      float var = var_smem[0] / D;                                          ║
║      float std = sqrt(var + eps);                                           ║
║                                                                              ║
║      // Phase 3: Normalize and output                                        ║
║      for (int i = 0; i < D_blk; i++) {                                       ║
║          int idx = bid * D + i * blockDim.x + tid;                            ║
║          float x_norm = (x[idx] - mean) / std;                               ║
║          y[idx] = gamma[idx % D] * x_norm + beta[idx % D];                  ║
║      }                                                                       ║
║  }                                                                           ║
║                                                                              ║
╚══════════════════════════════════════════════════════════════════════════════╝
""")

# =============================================================================
# Benchmark and Tests
# =============================================================================

def benchmark():
    """Benchmark forward and backward passes."""
    print("\n" + "="*70)
    print("BENCHMARKING LAYER NORMALIZATION")
    print("="*70)

    # Test different shapes
    shapes = [
        (32, 128, 256),    # Small
        (64, 128, 512),    # Medium (BERT-base hidden)
        (32, 512, 768),    # BERT-base
        (16, 512, 1024),   # Larger
    ]

    results = []

    for B, T, D in shapes:
        print(f"\nShape: (B={B}, T={T}, D={D})")
        print("-" * 40)

        # Create random inputs
        np.random.seed(42)
        x = np.random.randn(B, T, D).astype(np.float64)
        gamma = np.random.randn(D).astype(np.float64)
        beta = np.random.randn(D).astype(np.float64)
        dy = np.random.randn(B, T, D).astype(np.float64)

        # Forward benchmark
        n_iters = 100
        times = []
        for _ in range(n_iters):
            start = time.perf_counter()
            y, cache = layer_norm_forward(x, gamma, beta)
            end = time.perf_counter()
            times.append((end - start) * 1000)  # ms

        fwd_time = np.mean(times)
        fwd_std = np.std(times)

        # Backward benchmark
        times = []
        for _ in range(n_iters):
            start = time.perf_counter()
            dx, d_gamma, d_beta = layer_norm_backward(dy, cache)
            end = time.perf_counter()
            times.append((end - start) * 1000)  # ms

        bwd_time = np.mean(times)
        bwd_std = np.std(times)

        # Throughput
        elements = B * T * D
        fwd_throughput = elements / (fwd_time / 1000) / 1e9  # GB/s
        bwd_throughput = elements / (bwd_time / 1000) / 1e9

        print(f"Forward:  {fwd_time:.3f} ± {fwd_std:.3f} ms ({fwd_throughput:.1f} GB/s)")
        print(f"Backward: {bwd_time:.3f} ± {bwd_std:.3f} ms ({bwd_throughput:.1f} GB/s)")
        print(f"Total:    {fwd_time + bwd_time:.3f} ms")

        results.append({
            'shape': (B, T, D),
            'fwd_time': fwd_time,
            'bwd_time': bwd_time,
            'elements': elements
        })

    return results

def run_gradient_checks():
    """Run gradient checks on various shapes."""
    print("\n" + "="*70)
    print("RUNNING GRADIENT CHECKS")
    print("="*70)

    shapes = [
        (2, 4, 8),     # Tiny
        (4, 8, 16),    # Small
        (8, 16, 32),   # Medium-small
    ]

    all_passed = True

    for B, T, D in shapes:
        print(f"\n{'='*60}")
        print(f"Testing shape (B={B}, T={T}, D={D})")
        print(f"{'='*60}")

        # Create model and inputs
        model = LayerNorm(D)
        np.random.seed(123)
        x = np.random.randn(B, T, D).astype(np.float64)
        dy = np.random.randn(B, T, D).astype(np.float64)

        # Run gradient check
        results = gradient_check(model, x, dy, verbose=True)

        for name, passed in results.items():
            status = "✓ PASS" if passed else "✗ FAIL"
            print(f"  {name:>8}: {status}")
            all_passed = all_passed and passed

    print("\n" + "="*70)
    if all_passed:
        print("ALL GRADIENT CHECKS PASSED ✓")
    else:
        print("SOME GRADIENT CHECKS FAILED ✗")
    print("="*70)

    return all_passed

# =============================================================================
# Main
# =============================================================================

if __name__ == "__main__":
    print("="*70)
    print("LAYER NORMALIZATION: FORWARD AND BACKWARD FROM SCRATCH")
    print("="*70)

    # Print complexity analysis
    analyze_complexity()

    # Explain GPU fusion
    explain_gpu_fusion()

    # Run gradient checks
    run_gradient_checks()

    # Benchmark
    benchmark()


Successfully wrote 39773 bytes to ./minimax-m2.7/backwards/layer_norm_numpy.py


The file has been written. Let me now run it to verify everything works correctly.

bash

cd ./minimax-m2.7/backwards && python layer_norm_numpy.py


======================================================================
LAYER NORMALIZATION: FORWARD AND BACKWARD FROM SCRATCH
======================================================================

╔══════════════════════════════════════════════════════════════════════════════╗
║                    COMPLEXITY ANALYSIS                                         ║
╠══════════════════════════════════════════════════════════════════════════════╣
║                                                                              ║
║  Input: x ∈ ℝ^(B×T×D)                                                        ║
║  Parameters: γ, β ∈ ℝ^D                                                      ║
║                                                                              ║
╠══════════════════════════════════════════════════════════════════════════════╣
║  FORWARD PASS                                                                ║
╠══════════════════════════════════════════════════════════════════════════════╣
║                                                                              ║
║  Operation              │ Work (FLOPs)      │ Memory                        ║
║  ────────────────────────────────────────────────────────────────────────── ║
║  mean (reduction)       │ O(BTD)            │ -                             ║
║  x - mean (broadcast)   │ O(BTD)            │ O(BTD)                        ║
║  var (reduction)        │ O(BTD)            │ -                             ║
║  sqrt(var + eps)        │ O(BT)             │ -                             ║
║  divide by std          │ O(BTD)            │ -                             ║
║  gamma * x_norm         │ O(BTD)            │ -                             ║
║  add beta               │ O(BTD)            │ O(BTD) output                 ║
║  ────────────────────────────────────────────────────────────────────────── ║
║  TOTAL                    │ 5×O(BTD)          │ O(BTD)                       ║
║                                                                              ║
╠══════════════════════════════════════════════════════════════════════════════╣
║  BACKWARD PASS                                                              ║
╠══════════════════════════════════════════════════════════════════════════════╣
║                                                                              ║
║  Operation              │ Work (FLOPs)      │ Memory                        ║
║  ────────────────────────────────────────────────────────────────────────── ║
║  d_gamma = sum(dy*x_norm)│ O(BTD)           │ O(D)                          ║
║  d_beta = sum(dy)       │ O(BTD)            │ O(D)                          ║
║  dz = dy * gamma        │ O(BTD)            │ O(BTD) (can be avoided)        ║
║  sum(dz)                │ O(BTD)            │ -                             ║
║  sum(dz * x_norm)       │ O(BTD)            │ -                             ║
║  dx computation         │ O(BTD)            │ O(BTD) output                 ║
║  ────────────────────────────────────────────────────────────────────────── ║
║  TOTAL                    │ 5×O(BTD)          │ O(BTD) + O(D)                ║
║                                                                              ║
╠══════════════════════════════════════════════════════════════════════════════╣
║  SUMMARY                                                                    ║
╠══════════════════════════════════════════════════════════════════════════════╣
║                                                                              ║
║  Time Complexity:  O(BTD) for both forward and backward                      ║
║  Space Complexity: O(BTD) for storing activations (during training)          ║
║                   O(BTD) during inference (no need to store)                 ║
║                                                                              ║
║  Cache efficiency: We store x_centered, x_norm, mean, std                    ║
║                   These are O(BTD) total, reused across all gradient comps   ║
║                                                                              ║
╚══════════════════════════════════════════════════════════════════════════════╝

╔══════════════════════════════════════════════════════════════════════════════╗
║               GPU KERNEL FUSION DESIGN                                       ║
╠══════════════════════════════════════════════════════════════════════════════╣
║                                                                              ║
║  CURRENT SEPARATE KERNEL APPROACH:                                           ║
║  ────────────────────────────────────                                        ║
║  Kernel 1: Compute mean (reduction over D)                                   ║
║  Kernel 2: Compute variance (reduction over D)                               ║
║  Kernel 3: Normalize (element-wise)                                          ║
║  Kernel 4: Scale and shift (element-wise)                                    ║
║  Kernel 5: Backward kernels (x, gamma, beta)                                 ║
║                                                                              ║
║  Issues with separate kernels:                                               ║
║  • Multiple kernel launches (overhead)                                       ║
║  • Data movement between global memory passes                                 ║
║  • Can't use persistent threads for reduction efficiency                      ║
║                                                                              ║
╠══════════════════════════════════════════════════════════════════════════════╣
║  FUSED KERNEL DESIGN (Forward):                                              ║
╠══════════════════════════════════════════════════════════════════════════════╣
║                                                                              ║
║  Grid: (B × T) blocks, each handling one (b,t) position                      ║
║  Block: 256-512 threads                                                      ║
║                                                                              ║
║  PHASE 1: Load and compute local sum                                         ║
║  ─────────────────────────────────                                           ║
║  • Each thread loads x[b,t,d] into shared memory                             ║
║  • Compute partial sum using warp-level reduction                            ║
║  • Single thread writes mean to __shared__                                    ║
║                                                                              ║
║  PHASE 2: Compute variance locally                                           ║
║  ─────────────────────────────────                                           ║
║  • Re-load x with loaded mean                                                ║
║  • Compute (x-mean)² and partial variance                                    ║
║  • Reduce to get variance                                                   ║
║                                                                              ║
║  PHASE 3: Normalize and output                                               ║
║  ─────────────────────────────────                                           ║
║  • All threads compute: y = gamma * (x-mean) / sqrt(var+eps) + beta          ║
║  • Write to output (fully coalesced)                                         ║
║                                                                              ║
║  MEMORY ACCESS PATTERN:                                                      ║
║  • Each block reads contiguous D elements (coalesced)                        ║
║  • Use shared memory for intermediate results                                ║
║  • Output writes are also coalesced                                          ║
║                                                                              ║
╠══════════════════════════════════════════════════════════════════════════════╣
║  FUSED KERNEL DESIGN (Backward):                                             ║
╠══════════════════════════════════════════════════════════════════════════════╣
║                                                                              ║
║  Key insight: dz = dy * gamma can be merged into the computation              ║
║                                                                              ║
║  Grid: (B × T) blocks                                                       ║
║  Block: 256-512 threads                                                      ║
║                                                                              ║
║  SHARED MEMORY STRUCTURE:                                                    ║
║  [x_norm_0, x_norm_1, ..., x_norm_{D-1}]                                     ║
║  [dz_0, dz_1, ..., dz_{D-1}]                                                 ║
║                                                                              ║
║  ALGORITHM:                                                                  ║
║  1. Load x_norm and compute local dz = dy * gamma                            ║
║  2. Reduce to get sum(dz) and sum(dz * x_norm)                               ║
║  3. Second pass to compute dx using the formula:                             ║
║     dx = (dz - mean(dz) - x_norm * mean(dz*x_norm)) / std                    ║
║                                                                              ║
║  REDUCTION OPTIMIZATIONS:                                                    ║
║  • Warp-level shuffle reductions (no shared memory needed)                    ║
║  • Block-level using shared memory with tree reduction                        ║
║  • Use block-level primitives for final reduction                            ║
║                                                                              ║
╠══════════════════════════════════════════════════════════════════════════════╣
║  BENEFITS OF FUSION:                                                          ║
╠══════════════════════════════════════════════════════════════════════════════╣
║                                                                              ║
║  ✓ Reduced kernel launch overhead (1 kernel vs 4-5)                           ║
║  ✓ Better memory bandwidth utilization (single read, single write)          ║
║  ✓ Improved cache locality (data stays in registers/shared mem)              ║
║  ✓ Only loads x once, computes mean and var from same data                   ║
║  ✓ Backward can reuse cached values from forward (if memory allows)          ║
║  ✓ Lower register pressure allows for larger block sizes                     ║
║                                                                              ║
╠══════════════════════════════════════════════════════════════════════════════╣
║  CUDA KERNEL SKETCH (Pseudo-code):                                            ║
╠══════════════════════════════════════════════════════════════════════════════╣
║                                                                              ║
║  __global__ void layer_norm_fwd(float* y, const float* x,                    ║
║                                  const float* gamma, const float* beta,      ║
║                                  int B, int T, int D, float eps) {           ║
║                                                                              ║
║      __shared__ float mean_smem[256]; // block-level mean                     ║
║      __shared__ float var_smem[256];  // block-level variance                 ║
║      __shared__ float std_smem[256]; // block-level std                      ║
║                                                                              ║
║      int tid = threadIdx.x;                                                  ║
║      int bid = blockIdx.x;                                                   ║
║      int D_blk = (D + blockDim.x - 1) / blockDim.x;                          ║
║                                                                              ║
║      // Phase 1: Load and compute mean                                        ║
║      float sum = 0.0;                                                        ║
║      for (int i = 0; i < D_blk; i++) {                                       ║
║          int idx = bid * D + i * blockDim.x + tid;                           ║
║          sum += x[idx];                                                      ║
║      }                                                                       ║
║      sum = warpReduceSum(sum);                                               ║
║      if (tid % 32 == 0) mean_smem[tid / 32] = sum;                           ║
║      __syncthreads();                                                        ║
║                                                                              ║
║      float mean = mean_smem[0] / D;                                         ║
║                                                                              ║
║      // Phase 2: Compute variance                                            ║
║      sum = 0.0;                                                             ║
║      for (int i = 0; i < D_blk; i++) {                                       ║
║          int idx = bid * D + i * blockDim.x + tid;                           ║
║          float diff = x[idx] - mean;                                         ║
║          sum += diff * diff;                                                 ║
║      }                                                                       ║
║      sum = warpReduceSum(sum);                                               ║
║      if (tid % 32 == 0) var_smem[tid / 32] = sum;                           ║
║      __syncthreads();                                                       ║
║                                                                              ║
║      float var = var_smem[0] / D;                                          ║
║      float std = sqrt(var + eps);                                           ║
║                                                                              ║
║      // Phase 3: Normalize and output                                        ║
║      for (int i = 0; i < D_blk; i++) {                                       ║
║          int idx = bid * D + i * blockDim.x + tid;                            ║
║          float x_norm = (x[idx] - mean) / std;                               ║
║          y[idx] = gamma[idx % D] * x_norm + beta[idx % D];                  ║
║      }                                                                       ║
║  }                                                                           ║
║                                                                              ║
╚══════════════════════════════════════════════════════════════════════════════╝

======================================================================
RUNNING GRADIENT CHECKS
======================================================================

============================================================
Testing shape (B=2, T=4, D=8)
============================================================

============================================================
GRADIENT CHECK: gamma
============================================================
Traceback (most recent call last):
  File "./minimax-m2.7/backwards/layer_norm_numpy.py", line 920, in <module>
    run_gradient_checks()
  File "./minimax-m2.7/backwards/layer_norm_numpy.py", line 887, in run_gradient_checks
    results = gradient_check(model, x, dy, verbose=True)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "./minimax-m2.7/backwards/layer_norm_numpy.py", line 441, in gradient_check
    dx, d_gamma_analytical, _ = model.backward(dy, cache)
    ^^^^^^^^^^^^^^^^^^^^^^^^^
ValueError: not enough values to unpack (expected 3, got 2)

Command exited with code 1


There's a bug in the backward function - it's returning dx, d_gamma, d_beta but the gradient_check function expects the order to be (dx, d_gamma, d_beta). Let me check.

Actually looking at the code:
- `layer_norm_backward` returns `dx, d_gamma, d_beta`
- The model.backward calls `layer_norm_backward` and assigns `dx, self.d_gamma, self.d_beta = layer_norm_backward(dy, cache)`

So that's fine. But the gradient_check function has:
```python
dx, d_gamma_analytical, _ = model.backward(dy, cache)
```

Wait, but model.backward only returns dx. Let me check the LayerNorm class:

```python
def backward(self, dy: np.ndarray, cache: Dict) -> np.ndarray:
    ...
    return dx
```

The model.backward method only returns dx! But in gradient_check, we try to unpack 3 values. That's the bug.

I need to modify the model.backward to return all three, or create a different function. Let me fix this.

I need to fix a bug in the `backward` method - it only returns `dx` but gradient check expects the full `(dx, d_gamma, d_beta)`. Let me fix this:

edit

    def backward(self, dy: np.ndarray, cache: Dict) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
        """
        Backward pass.

        Args:
            dy: Upstream gradient of shape (B, T, D)
            cache: Forward pass cache

        Returns:
            dx: Gradient w.r.t. input x
            d_gamma: Gradient w.r.t. gamma
            d_beta: Gradient w.r.t. beta
        """
        dx, d_gamma, d_beta = layer_norm_backward(dy, cache)
        self.d_gamma = d_gamma
        self.d_beta = d_beta
        return dx, d_gamma, d_beta

    def backward(self, dy: np.ndarray, cache: Dict) -> np.ndarray:
        """
        Backward pass.

        Args:
            dy: Upstream gradient of shape (B, T, D)
            cache: Forward pass cache

        Returns:
            dx: Gradient w.r.t. input x
        """
        dx, self.d_gamma, self.d_beta = layer_norm_backward(dy, cache)
        return dx


Successfully replaced 1 block(s) in ./minimax-m2.7/backwards/layer_norm_numpy.py.

     ...
 261         assert x.shape[-1] == self.normalized_shape, \
 262             f"Expected last dimension {self.normalized_shape}, got {x.shape[-1]}"
 263         return layer_norm_forward(x, self.gamma, self.beta, self.eps)
 264
-265     def backward(self, dy: np.ndarray, cache: Dict) -> np.ndarray:
+265     def backward(self, dy: np.ndarray, cache: Dict) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
 266         """
 267         Backward pass.
 268
 269         Args:
     ...
 271             cache: Forward pass cache
 272
 273         Returns:
 274             dx: Gradient w.r.t. input x
+275             d_gamma: Gradient w.r.t. gamma
+276             d_beta: Gradient w.r.t. beta
 275         """
-276         dx, self.d_gamma, self.d_beta = layer_norm_backward(dy, cache)
-277         return dx
+278         dx, d_gamma, d_beta = layer_norm_backward(dy, cache)
+279         self.d_gamma = d_gamma
+280         self.d_beta = d_beta
+281         return dx, d_gamma, d_beta
 278
 279     def parameters(self) -> Tuple[np.ndarray, np.ndarray]:
 280         """Return (gamma, beta)."""
 281         return self.gamma, self.beta
     ...


Now I need to fix the gradient_check function to handle the updated return signature.

    model.gamma = orig_gamma.copy()
    y, cache = model.forward(x)
    dx_analytical, d_gamma_analytical, _ = model.backward(dy, cache)

    model.gamma = orig_gamma.copy()
    y, cache = model.forward(x)
    dx, d_gamma_analytical, _ = model.backward(dy, cache)


     ...
 441         print("="*60)
 442
 443     model.gamma = orig_gamma.copy()
 444     y, cache = model.forward(x)
-445     dx, d_gamma_analytical, _ = model.backward(dy, cache)
+445     dx_analytical, d_gamma_analytical, _ = model.backward(dy, cache)
 446
 447     d_gamma_numerical = compute_numerical_gradient(
 448         func_with_param, x, orig_gamma.copy(), 'gamma', dy
 449     )
     ...


Let me also check the other gradient check section for beta.


    diff = np.abs(d_gamma_analytical - d_gamma_numerical)
    tolerance = atol + rtol * np.abs(d_gamma_numerical)
    passed = np.all(diff <= tolerance)

    if verbose:
        print(f"Analytical gradient shape: {d_gamma_analytical.shape}")
        print(f"Numerical gradient shape: {d_gamma_numerical.shape}")
        print(f"Max absolute difference: {np.max(diff):.2e}")
        print(f"Max relative tolerance: {np.max(tolerance):.2e}")
        print(f"Mean analytical: {np.mean(np.abs(d_gamma_analytical)):.6e}")
        print(f"Mean numerical: {np.mean(np.abs(d_gamma_numerical)):.6e}")
        print(f"\nGradient check: {'PASSED ✓' if passed else 'FAILED ✗'}")

    results['gamma'] = passed

    # -------------------------------------------------------------------------
    # Check gradient w.r.t. beta
    # -------------------------------------------------------------------------
    if verbose:
        print("\n" + "="*60)
        print("GRADIENT CHECK: beta")
        print("="*60)

    model.beta = orig_beta.copy()
    y, cache = model.forward(x)
    dx, _, d_beta_analytical = model.backward(dy, cache)

    d_beta_numerical = compute_numerical_gradient(
        func_with_param, x, orig_beta.copy(), 'beta', dy
    )

    # Reset beta
    model.beta = orig_beta.copy()

    diff = np.abs(d_beta_analytical - d_beta_numerical)
    tolerance = atol + rtol * np.abs(d_beta_numerical)
    passed = np.all(diff <= tolerance)

    if verbose:
        print(f"Analytical gradient shape: {d_beta_analytical.shape}")
        print(f"Numerical gradient shape: {d_beta_numerical.shape}")
        print(f"Max absolute difference: {np.max(diff):.2e}")
        print(f"Max relative tolerance: {np.max(tolerance):.2e}")
        print(f"Mean analytical: {np.mean(np.abs(d_beta_analytical)):.6e}")
        print(f"Mean numerical: {np.mean(np.abs(d_beta_numerical)):.6e}")
        print(f"\nGradient check: {'PASSED ✓' if passed else 'FAILED ✗'}")

    results['beta'] = passed

    # -------------------------------------------------------------------------
    # Check gradient w.r.t. x
    # -------------------------------------------------------------------------
    if verbose:
        print("\n" + "="*60)
        print("GRADIENT CHECK: x (input)")
        print("="*60)

    # Need to recompute since model was modified
    model.gamma = orig_gamma.copy()
    model.beta = orig_beta.copy()

[413 more lines in file. Use offset=515 to continue.]


I see the issue. On line 471 and 492, I have `dx, _, d_beta_analytical = model.backward(dy, cache)` but with the new signature, the return is `dx, d_gamma, d_beta`. Let me fix these:

    model.beta = orig_beta.copy()
    y, cache = model.forward(x)
    _, _, d_beta_analytical = model.backward(dy, cache)

    model.beta = orig_beta.copy()
    y, cache = model.forward(x)
    dx, _, d_beta_analytical = model.backward(dy, cache)

    model.gamma = orig_gamma.copy()
    model.beta = orig_beta.copy()
    y, cache = model.forward(x)
    dx_analytical, _, _ = model.backward(dy, cache)


Successfully replaced 2 block(s) in ./minimax-m2.7/backwards/layer_norm_numpy.py.

     ...
 476         print("="*60)
 477
 478     model.beta = orig_beta.copy()
 479     y, cache = model.forward(x)
-480     dx, _, d_beta_analytical = model.backward(dy, cache)
+480     _, _, d_beta_analytical = model.backward(dy, cache)
 481
 482     d_beta_numerical = compute_numerical_gradient(
 483         func_with_param, x, orig_beta.copy(), 'beta', dy
 484     )
     ...


======================================================================
LAYER NORMALIZATION: FORWARD AND BACKWARD FROM SCRATCH
======================================================================

╔══════════════════════════════════════════════════════════════════════════════╗
║                    COMPLEXITY ANALYSIS                                         ║
╠══════════════════════════════════════════════════════════════════════════════╣
║                                                                              ║
║  Input: x ∈ ℝ^(B×T×D)                                                        ║
║  Parameters: γ, β ∈ ℝ^D                                                      ║
║                                                                              ║
╠══════════════════════════════════════════════════════════════════════════════╣
║  FORWARD PASS                                                                ║
╠══════════════════════════════════════════════════════════════════════════════╣
║                                                                              ║
║  Operation              │ Work (FLOPs)      │ Memory                        ║
║  ────────────────────────────────────────────────────────────────────────── ║
║  mean (reduction)       │ O(BTD)            │ -                             ║
║  x - mean (broadcast)   │ O(BTD)            │ O(BTD)                        ║
║  var (reduction)        │ O(BTD)            │ -                             ║
║  sqrt(var + eps)        │ O(BT)             │ -                             ║
║  divide by std          │ O(BTD)            │ -                             ║
║  gamma * x_norm         │ O(BTD)            │ -                             ║
║  add beta               │ O(BTD)            │ O(BTD) output                 ║
║  ────────────────────────────────────────────────────────────────────────── ║
║  TOTAL                    │ 5×O(BTD)          │ O(BTD)                       ║
║                                                                              ║
╠══════════════════════════════════════════════════════════════════════════════╣
║  BACKWARD PASS                                                              ║
╠══════════════════════════════════════════════════════════════════════════════╣
║                                                                              ║
║  Operation              │ Work (FLOPs)      │ Memory                        ║
║  ────────────────────────────────────────────────────────────────────────── ║
║  d_gamma = sum(dy*x_norm)│ O(BTD)           │ O(D)                          ║
║  d_beta = sum(dy)       │ O(BTD)            │ O(D)                          ║
║  dz = dy * gamma        │ O(BTD)            │ O(BTD) (can be avoided)        ║
║  sum(dz)                │ O(BTD)            │ -                             ║
║  sum(dz * x_norm)       │ O(BTD)            │ -                             ║
║  dx computation         │ O(BTD)            │ O(BTD) output                 ║
║  ────────────────────────────────────────────────────────────────────────── ║
║  TOTAL                    │ 5×O(BTD)          │ O(BTD) + O(D)                ║
║                                                                              ║
╠══════════════════════════════════════════════════════════════════════════════╣
║  SUMMARY                                                                    ║
╠══════════════════════════════════════════════════════════════════════════════╣
║                                                                              ║
║  Time Complexity:  O(BTD) for both forward and backward                      ║
║  Space Complexity: O(BTD) for storing activations (during training)          ║
║                   O(BTD) during inference (no need to store)                 ║
║                                                                              ║
║  Cache efficiency: We store x_centered, x_norm, mean, std                    ║
║                   These are O(BTD) total, reused across all gradient comps   ║
║                                                                              ║
╚══════════════════════════════════════════════════════════════════════════════╝

╔══════════════════════════════════════════════════════════════════════════════╗
║               GPU KERNEL FUSION DESIGN                                       ║
╠══════════════════════════════════════════════════════════════════════════════╣
║                                                                              ║
║  CURRENT SEPARATE KERNEL APPROACH:                                           ║
║  ────────────────────────────────────                                        ║
║  Kernel 1: Compute mean (reduction over D)                                   ║
║  Kernel 2: Compute variance (reduction over D)                               ║
║  Kernel 3: Normalize (element-wise)                                          ║
║  Kernel 4: Scale and shift (element-wise)                                    ║
║  Kernel 5: Backward kernels (x, gamma, beta)                                 ║
║                                                                              ║
║  Issues with separate kernels:                                               ║
║  • Multiple kernel launches (overhead)                                       ║
║  • Data movement between global memory passes                                 ║
║  • Can't use persistent threads for reduction efficiency                      ║
║                                                                              ║
╠══════════════════════════════════════════════════════════════════════════════╣
║  FUSED KERNEL DESIGN (Forward):                                              ║
╠══════════════════════════════════════════════════════════════════════════════╣
║                                                                              ║
║  Grid: (B × T) blocks, each handling one (b,t) position                      ║
║  Block: 256-512 threads                                                      ║
║                                                                              ║
║  PHASE 1: Load and compute local sum                                         ║
║  ─────────────────────────────────                                           ║
║  • Each thread loads x[b,t,d] into shared memory                             ║
║  • Compute partial sum using warp-level reduction                            ║
║  • Single thread writes mean to __shared__                                    ║
║                                                                              ║
║  PHASE 2: Compute variance locally                                           ║
║  ─────────────────────────────────                                           ║
║  • Re-load x with loaded mean                                                ║
║  • Compute (x-mean)² and partial variance                                    ║
║  • Reduce to get variance                                                   ║
║                                                                              ║
║  PHASE 3: Normalize and output                                               ║
║  ─────────────────────────────────                                           ║
║  • All threads compute: y = gamma * (x-mean) / sqrt(var+eps) + beta          ║
║  • Write to output (fully coalesced)                                         ║
║                                                                              ║
║  MEMORY ACCESS PATTERN:                                                      ║
║  • Each block reads contiguous D elements (coalesced)                        ║
║  • Use shared memory for intermediate results                                ║
║  • Output writes are also coalesced                                          ║
║                                                                              ║
╠══════════════════════════════════════════════════════════════════════════════╣
║  FUSED KERNEL DESIGN (Backward):                                             ║
╠══════════════════════════════════════════════════════════════════════════════╣
║                                                                              ║
║  Key insight: dz = dy * gamma can be merged into the computation              ║
║                                                                              ║
║  Grid: (B × T) blocks                                                       ║
║  Block: 256-512 threads                                                      ║
║                                                                              ║
║  SHARED MEMORY STRUCTURE:                                                    ║
║  [x_norm_0, x_norm_1, ..., x_norm_{D-1}]                                     ║
║  [dz_0, dz_1, ..., dz_{D-1}]                                                 ║
║                                                                              ║
║  ALGORITHM:                                                                  ║
║  1. Load x_norm and compute local dz = dy * gamma                            ║
║  2. Reduce to get sum(dz) and sum(dz * x_norm)                               ║
║  3. Second pass to compute dx using the formula:                             ║
║     dx = (dz - mean(dz) - x_norm * mean(dz*x_norm)) / std                    ║
║                                                                              ║
║  REDUCTION OPTIMIZATIONS:                                                    ║
║  • Warp-level shuffle reductions (no shared memory needed)                    ║
║  • Block-level using shared memory with tree reduction                        ║
║  • Use block-level primitives for final reduction                            ║
║                                                                              ║
╠══════════════════════════════════════════════════════════════════════════════╣
║  BENEFITS OF FUSION:                                                          ║
╠══════════════════════════════════════════════════════════════════════════════╣
║                                                                              ║
║  ✓ Reduced kernel launch overhead (1 kernel vs 4-5)                           ║
║  ✓ Better memory bandwidth utilization (single read, single write)          ║
║  ✓ Improved cache locality (data stays in registers/shared mem)              ║
║  ✓ Only loads x once, computes mean and var from same data                   ║
║  ✓ Backward can reuse cached values from forward (if memory allows)          ║
║  ✓ Lower register pressure allows for larger block sizes                     ║
║                                                                              ║
╠══════════════════════════════════════════════════════════════════════════════╣
║  CUDA KERNEL SKETCH (Pseudo-code):                                            ║
╠══════════════════════════════════════════════════════════════════════════════╣
║                                                                              ║
║  __global__ void layer_norm_fwd(float* y, const float* x,                    ║
║                                  const float* gamma, const float* beta,      ║
║                                  int B, int T, int D, float eps) {           ║
║                                                                              ║
║      __shared__ float mean_smem[256]; // block-level mean                     ║
║      __shared__ float var_smem[256];  // block-level variance                 ║
║      __shared__ float std_smem[256]; // block-level std                      ║
║                                                                              ║
║      int tid = threadIdx.x;                                                  ║
║      int bid = blockIdx.x;                                                   ║
║      int D_blk = (D + blockDim.x - 1) / blockDim.x;                          ║
║                                                                              ║
║      // Phase 1: Load and compute mean                                        ║
║      float sum = 0.0;                                                        ║
║      for (int i = 0; i < D_blk; i++) {                                       ║
║          int idx = bid * D + i * blockDim.x + tid;                           ║
║          sum += x[idx];                                                      ║
║      }                                                                       ║
║      sum = warpReduceSum(sum);                                               ║
║      if (tid % 32 == 0) mean_smem[tid / 32] = sum;                           ║
║      __syncthreads();                                                        ║
║                                                                              ║
║      float mean = mean_smem[0] / D;                                         ║
║                                                                              ║
║      // Phase 2: Compute variance                                            ║
║      sum = 0.0;                                                             ║
║      for (int i = 0; i < D_blk; i++) {                                       ║
║          int idx = bid * D + i * blockDim.x + tid;                           ║
║          float diff = x[idx] - mean;                                         ║
║          sum += diff * diff;                                                 ║
║      }                                                                       ║
║      sum = warpReduceSum(sum);                                               ║
║      if (tid % 32 == 0) var_smem[tid / 32] = sum;                           ║
║      __syncthreads();                                                       ║
║                                                                              ║
║      float var = var_smem[0] / D;                                          ║
║      float std = sqrt(var + eps);                                           ║
║                                                                              ║
║      // Phase 3: Normalize and output                                        ║
║      for (int i = 0; i < D_blk; i++) {                                       ║
║          int idx = bid * D + i * blockDim.x + tid;                            ║
║          float x_norm = (x[idx] - mean) / std;                               ║
║          y[idx] = gamma[idx % D] * x_norm + beta[idx % D];                  ║
║      }                                                                       ║
║  }                                                                           ║
║                                                                              ║
╚══════════════════════════════════════════════════════════════════════════════╝

======================================================================
RUNNING GRADIENT CHECKS
======================================================================

============================================================
Testing shape (B=2, T=4, D=8)
============================================================

============================================================
GRADIENT CHECK: gamma
============================================================
Analytical gradient shape: (8,)
Numerical gradient shape: (8,)
Max absolute difference: 3.67e+00
Max relative tolerance: 1.00e-05
Mean analytical: 2.081636e+00
Mean numerical: 0.000000e+00

Gradient check: FAILED ✗

============================================================
GRADIENT CHECK: beta
============================================================
Analytical gradient shape: (8,)
Numerical gradient shape: (8,)
Max absolute difference: 4.02e+00
Max relative tolerance: 1.00e-05
Mean analytical: 1.957784e+00
Mean numerical: 0.000000e+00

Gradient check: FAILED ✗

============================================================
GRADIENT CHECK: x (input)
============================================================
  Computing numerical gradient for x (64 elements)...
Analytical gradient shape: (2, 4, 8)
Numerical gradient shape: (2, 4, 8)
Max absolute difference: 2.99e+00
Max relative tolerance: 1.00e-05
Mean analytical: 7.927930e-01
Mean numerical: 0.000000e+00

Gradient check: FAILED ✗
     gamma: ✗ FAIL
      beta: ✗ FAIL
         x: ✗ FAIL

============================================================
Testing shape (B=4, T=8, D=16)
============================================================

============================================================
GRADIENT CHECK: gamma
============================================================
Analytical gradient shape: (16,)
Numerical gradient shape: (16,)
Max absolute difference: 1.19e+01
Max relative tolerance: 1.00e-05
Mean analytical: 4.686865e+00
Mean numerical: 0.000000e+00

Gradient check: FAILED ✗

============================================================
GRADIENT CHECK: beta
============================================================
Analytical gradient shape: (16,)
Numerical gradient shape: (16,)
Max absolute difference: 8.73e+00
Max relative tolerance: 1.00e-05
Mean analytical: 3.191375e+00
Mean numerical: 0.000000e+00

Gradient check: FAILED ✗

============================================================
GRADIENT CHECK: x (input)
============================================================
  Computing numerical gradient for x (512 elements)...
Analytical gradient shape: (4, 8, 16)
Numerical gradient shape: (4, 8, 16)
Max absolute difference: 4.19e+00
Max relative tolerance: 1.00e-05
Mean analytical: 7.809152e-01
Mean numerical: 0.000000e+00

Gradient check: FAILED ✗
     gamma: ✗ FAIL
      beta: ✗ FAIL
         x: ✗ FAIL

============================================================
Testing shape (B=8, T=16, D=32)
============================================================

============================================================
GRADIENT CHECK: gamma
============================================================
Analytical gradient shape: (32,)
Numerical gradient shape: (32,)
Max absolute difference: 2.34e+01
Max relative tolerance: 1.00e-05
Mean analytical: 7.617486e+00
Mean numerical: 0.000000e+00

Gradient check: FAILED ✗

============================================================
GRADIENT CHECK: beta
============================================================
Analytical gradient shape: (32,)
Numerical gradient shape: (32,)
Max absolute difference: 2.62e+01
Max relative tolerance: 1.00e-05
Mean analytical: 9.277827e+00
Mean numerical: 0.000000e+00

Gradient check: FAILED ✗

============================================================
GRADIENT CHECK: x (input)
============================================================
  Computing numerical gradient for x (4096 elements)...
Analytical gradient shape: (8, 16, 32)
Numerical gradient shape: (8, 16, 32)
Max absolute difference: 4.12e+00
Max relative tolerance: 1.00e-05
Mean analytical: 8.156518e-01
Mean numerical: 0.000000e+00

Gradient check: FAILED ✗
     gamma: ✗ FAIL
      beta: ✗ FAIL
         x: ✗ FAIL

======================================================================
SOME GRADIENT CHECKS FAILED ✗
======================================================================

======================================================================
BENCHMARKING LAYER NORMALIZATION
======================================================================

Shape: (B=32, T=128, D=256)
----------------------------------------
Forward:  6.626 ± 0.649 ms (0.2 GB/s)
Backward: 8.112 ± 0.993 ms (0.1 GB/s)
Total:    14.738 ms

Shape: (B=64, T=128, D=512)
----------------------------------------
Forward:  26.815 ± 2.229 ms (0.2 GB/s)
Backward: 34.261 ± 2.739 ms (0.1 GB/s)
Total:    61.076 ms

Shape: (B=32, T=512, D=768)
----------------------------------------
Forward:  84.485 ± 6.490 ms (0.1 GB/s)
Backward: 102.953 ± 7.625 ms (0.1 GB/s)
Total:    187.437 ms

Shape: (B=16, T=512, D=1024)
----------------------------------------
Forward:  53.035 ± 3.852 ms (0.2 GB/s)
Backward: 69.202 ± 3.853 ms (0.1 GB/s)
Total:    122.237 ms


The gradient checks are failing. The numerical gradients are all zeros, which means the numerical gradient computation function has a bug. Let me look at the `compute_numerical_gradient` function more carefully.

The issue is that the function takes `param` but the way it's being used is wrong. Let me trace through:

1. In `gradient_check`, we call:
   ```python
   d_gamma_numerical = compute_numerical_gradient(
       func_with_param, x, orig_gamma.copy(), 'gamma', dy
   )
   ```

2. In `compute_numerical_gradient`:
   ```python
   param_flat = param.flatten()
   for i in range(len(param_flat)):
       original = param_flat[i]
       param_flat[i] = original + h
       y_plus, _ = func_with_param(x, param)  # But param here is a COPY of orig_gamma!
       ...
   ```

The problem is that `func_with_param` takes the model and modifies `model.gamma` or `model.beta` to match the provided `param`. But when I pass `orig_gamma.copy()` to `compute_numerical_gradient`, and inside the function I'm modifying `param_flat` which is a copy, the changes to `param` don't affect what `func_with_param` sees because `func_with_param` reads from `model.gamma` or `model.beta`, not from the passed `param`.

Actually wait, let me re-read `func_with_param`:

```python
def func_with_param(x_input, param):
    # Temporarily set the parameter
    if param is model.gamma:
        model.gamma = param
    else:
        model.beta = param
    y, cache = model.forward(x_input)
    # Restore original
    if param is model.gamma:
        model.gamma = orig_gamma
    else:
        model.beta = orig_beta
    return y, cache
```

The core issue is that when I pass `orig_gamma.copy()`, the identity check `param is model.gamma` fails since `orig_gamma.copy()` creates a new object. This means the function never actually updates `model.gamma` during the gradient computation, so the forward pass always sees the original gamma values regardless of the perturbations being tested.

Additionally, the restore logic uses `orig_gamma` directly instead of `orig_beta`, creating a bug. I should either use direct assignment to model parameters and restore using saved copies, or better yet, pass the original parameter directly and let the function handle setting it via `model.gamma = param` before computing the loss.

The core issue is that identity checks fail when we're comparing against copies of the original arrays. I need to restructure this so that the function receives the actual model parameters and can set them directly before the forward pass, then restore them afterward.

The gradient check is failing because the numerical gradient function has a bug. Let me fix it:


    x: np.ndarray,
    param: np.ndarray,
    param_name: str,
    dy: np.ndarray,
    h: float = 1e-5,
    cache_base: Optional[Dict] = None
) -> np.ndarray:
    """
    Compute numerical gradient using finite differences.

    Uses central difference: df/dx ≈ (f(x+h) - f(x-h)) / (2h)

    Args:
        func: Function that takes (x, param) and returns (y, cache)
        x: Input tensor
        param: Parameter array to compute gradient for
        param_name: Name of parameter ('gamma' or 'beta')
        dy: Upstream gradient
        h: Step size for finite differences
        cache_base: Base cache if modifying param

    Returns:
        Numerical gradient of same shape as param
    """
    num_grad = np.zeros_like(param)

    # Flatten for easier iteration
    param_flat = param.flatten()
    num_grad_flat = num_grad.flatten()

    for i in range(len(param_flat)):
        # Save original value
        original = param_flat[i]

        # f(x, param + h)
        param_flat[i] = original + h
        y_plus, _ = func(x, param)
        loss_plus = np.sum(y_plus * dy)

        # f(x, param - h)
        param_flat[i] = original - h
        y_minus, _ = func(x, param)
        loss_minus = np.sum(y_minus * dy)

        # Central difference
        num_grad_flat[i] = (loss_plus - loss_minus) / (2 * h)

        # Restore original
        param_flat[i] = original

    return num_grad_flat.reshape(param.shape)

def compute_numerical_gradient_wrt_x(
    model: LayerNorm,
    x: np.ndarray,
    dy: np.ndarray,
    h: float = 1e-5
) -> np.ndarray:
    """

[570 more lines in file. Use offset=358 to continue.]


I see the issue now. The `param_flat[i] = original + h` doesn't work because `param` and `param_flat` share the same memory. When I do `param_flat[i] = original + h`, it's also modifying `param`. But the problem is that `y_plus, _ = func(x, param)` is called after modifying the array, and the array has been modified.

Wait, that's exactly what we want! When we set `param_flat[i] = original + h`, we're modifying the underlying array, so `func(x, param)` should see the modified parameter.

But there's a subtle issue: after the first iteration, the original value has been restored to `param_flat[i]`. Let me trace through:

1. First iteration (i=0):   - original = param_flat[0] (say value v)   - param_flat[0] = v + h   - y_plus = f(x, param)  -- param[0] = v + h ✓   - param_flat[0] = v - h   - y_minus = f(x, param) -- param[0] = v - h ✓   - param_flat[0] = v  -- restored ✓

2. Second iteration (i=1):
   - original = param_flat[1] (say value w)
   - param_flat[1] = w + h
   - y_plus = f(x, param)  -- param[1] = w + h ✓
   - param_flat[1] = w - h
   - y_minus = f(x, param) -- param[1] = w - h ✓
   - param_flat[1] = w  -- restored ✓

This looks correct to me. The issue must be somewhere else.

Wait, let me look at what `param` actually is. It's passed as an argument, but the `param_flat` is just a view into `param` (since we flatten it).

Since they share memory, modifying `param_flat[i]` changes `param[i]`. So when I do `param_flat[i] = original + h`, the underlying `param` array gets modified, which is what the function should see.

The problem might be in how `func_with_param` works. I need to check that function again to see if there's something preventing the modifications from being applied correctly.

When `func_with_param` is called with `param`, it modifies `model.gamma` or `model.beta` to match `param`, then calls `model.forward(x)`. After that, it restores the original values and returns the output and cache. But if the modifications aren't actually affecting the model's parameters before `forward` is called, then the loss calculations won't reflect the intended parameter changes.

Let me trace through what happens when `compute_numerical_gradient` is invoked with `func_with_param` as the function and `orig_gamma` as the parameter. The function flattens the parameter, then for each element i, it temporarily modifies that element to `original + h` before calling `func_with_param(x, param)` with the flattened parameter array. The issue is whether these modifications to `param_flat` actually propagate to the model's internal parameters through the logic in `func_with_param`.

When `func_with_param` receives `param`, it reassigns `model.gamma` to point to that same array. After `model.forward(x)` executes and `model.gamma` gets reset to its original values, the `param` variable still references the modified array. At the end of the iteration, the code restores `param_flat[i]` back to its original value, but since `param` is now pointing to that same array, this creates a mismatch between what the model has and what the parameter variable holds.

The key insight is that modifying `param_flat[i]` directly changes the underlying array that both `param_flat` and `param` reference. So when `func_with_param` runs with `param`, it's modifying and restoring `model.gamma` correctly, but the side effect is that `param` itself gets modified in the process—though it gets restored to its original state by the end of each iteration.

In the next iteration, we retrieve the actual value from `param_flat[1]` (which remains unchanged since we only modified index 0), then modify index 1 and compute the function. After restoring, we move to index 2 and repeat the pattern.

For each subsequent iteration, we load the current element's value, update that element by adding h, call the function with the modified parameter, subtract h to get the minus version, call the function again, compute the gradient estimate via finite differences, then restore the original value. This continues until we've processed all elements.

I notice a critical bug in the restoration step: the code unconditionally sets `param_flat[i] = original` after computing y_minus, which overwrites the v - h value that was just used in the function call. This means we're not actually preserving the state we had before the iteration started.

Actually, wait—if the last modification was supposed to be the restore operation, then this is correct. We need to put back the original value, not leave it at v - h. The order of operations matters: we first set it to v + h and computed the function, then set it to v - h and computed the function again, and now we need to restore it to v. The finite difference calculation itself happens between these modifications, so it doesn't affect the state we're trying to maintain.

Let me verify the mechanics: when we modify `param_flat`, we're changing the array that `param` references, which in turn affects what `model.gamma` sees when it's set to point to `param`. The key is that `model.gamma` and `param` share the same underlying memory, so any modifications to the flattened view propagate correctly through the model parameter.

But actually, I should double-check whether this is actually working as expected. The issue might be that `np.zeros_like(param)` is creating a copy with the same dtype, and if `param` is float64, then `num_grad` should also be float64.

Let me trace through what's happening more carefully. When I modify `param_flat[i]`, I'm changing the actual array elements that both `param` and `model.gamma` reference. The function call should see these modifications. But maybe the problem is that the function isn't actually using the modified values, or perhaps the numerical gradient is genuinely zero for some reason. I should run a simpler test case to see what's actually happening.

bash

cd ./minimax-m2.7/backwards && python3 -c "
import numpy as np

def layer_norm_forward(x, gamma, beta, eps=1e-8):
    B, T, D = x.shape
    mean = np.mean(x, axis=-1, keepdims=True)
    x_centered = x - mean
    var = np.mean(x_centered ** 2, axis=-1, keepdims=True)
    std = np.sqrt(var + eps)
    x_norm = x_centered / std
    y = gamma * x_norm + beta
    return y, {'mean': mean, 'std': std, 'x_norm': x_norm, 'x_centered': x_centered}

def layer_norm_backward(dy, cache):
    x_norm = cache['x_norm']
    std = cache['std']
    B, T, D = x_norm.shape

    d_gamma = np.sum(dy * x_norm, axis=(0, 1))
    d_beta = np.sum(dy, axis=(0, 1))

    dz = dy * cache.get('gamma', np.ones(D))  # simplified
    sum_dz = np.sum(dz, axis=-1, keepdims=True)
    sum_dz_xnorm = np.sum(dz * x_norm, axis=-1, keepdims=True)
    dx = (dz - sum_dz / D - x_norm * sum_dz_xnorm / D) / std

    return dx, d_gamma, d_beta

# Test with simple values
np.random.seed(42)
B, T, D = 2, 4, 8
x = np.random.randn(B, T, D).astype(np.float64)
gamma = np.ones(D).astype(np.float64)  # Start with 1s
beta = np.zeros(D).astype(np.float64)
dy = np.random.randn(B, T, D).astype(np.float64)

# Manual gradient check for gamma[0]
h = 1e-5

# Analytical gradient
y, cache = layer_norm_forward(x, gamma, beta)
dx, d_gamma_analytical, _ = layer_norm_backward(dy, cache)
print(f'Analytical d_gamma: {d_gamma_analytical}')

# Numerical gradient for gamma[0] only
def loss_fn(x, g, b):
    y, _ = layer_norm_forward(x, g, b)
    return np.sum(y * dy)

loss_base = loss_fn(x, gamma, beta)
gamma_plus = gamma.copy()
gamma_plus[0] += h
loss_plus = loss_fn(x, gamma_plus, beta)
gamma_minus = gamma.copy()
gamma_minus[0] -= h
loss_minus = loss_fn(x, gamma_minus, beta)

d_gamma_num = (loss_plus - loss_minus) / (2 * h)
print(f'Numerical d_gamma[0]: {d_gamma_num}')
print(f'Difference: {abs(d_gamma_analytical[0] - d_gamma_num):.2e}')
"


Analytical d_gamma: [-0.20937132 -3.81744267 -0.006511   -0.04104874 -3.23976352  1.70935885
  4.55825163 -2.56074582]
Numerical d_gamma[0]: -0.20937132299003733
Difference: 1.87e-11