The analytical gradient for gamma[0] matches the numerical gradient! The issue is in how the gradient check is implemented. Let me look at the `func_with_param` function and the loop to understand the problem.


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_analytical, 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:

[468 more lines in file. Use offset=460 to continue.]


I see the issue now. The `compute_numerical_gradient` function is wrong. The `func_with_param` function is modifying and restoring `model.gamma` inside, but it's receiving `orig_gamma.copy()` as `param`, which is not the model's current gamma! Let me fix this.

Also, the function returns `y, cache` but the numerical gradient function is doing the right thing by calling `func(x, param)` - but the issue is that `param` passed to the numerical gradient is `orig_gamma.copy()`, not a reference to `model.gamma`.

Let me rewrite the gradient_check function to be cleaner and fix this bug:

I see the issue. The numerical gradient function has a bug - it's passing `orig_gamma.copy()` but the helper function has issues with reference handling. Let me fix this:


        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_analytical, 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)
    _, _, 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}")

[432 more lines in file. Use offset=496 to continue.]


Now I see the issue. The `gradient_check` function resets the model.gamma/model.beta AFTER computing the analytical gradient, but then calls `compute_numerical_gradient` which uses `func_with_param`. The `func_with_param` compares `param is model.gamma` but by that point the model.gamma has been reset.

Let me trace through:
1. `orig_gamma = model.gamma.copy()` - stores the original gamma
2. `model.gamma = orig_gamma.copy()` - sets model.gamma to the original
3. `y, cache = model.forward(x)` - forward pass
4. `dx_analytical, d_gamma_analytical, _ = model.backward(dy, cache)` - backward pass (analytical)
5. `d_gamma_numerical = compute_numerical_gradient(func_with_param, x, orig_gamma.copy(), 'gamma', dy)`:
   - In compute_numerical_gradient, param = orig_gamma.copy()
   - When we do `param is model.gamma`, we're comparing a NEW array (orig_gamma.copy()) to model.gamma
   - This will be False! So we go into the else branch and set model.beta = param
   - Then we do `y, cache = model.forward(x)` but model.gamma is still orig_gamma (not the perturbed one)!
   - This is the bug!

The issue is the identity check `param is model.gamma`. The numerical gradient function creates a copy of the parameter and passes that as `param`. When we do `param is model.gamma`, the identity check fails because `param` is a new array (copy), not the same object.

Let me rewrite the gradient check to fix this.

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
import copy

# =============================================================================
# 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) -> 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 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_gamma(
    x: np.ndarray,
    gamma: np.ndarray,
    beta: np.ndarray,
    dy: np.ndarray,
    h: float = 1e-5
) -> np.ndarray:
    """
    Compute numerical gradient for gamma using finite differences.

    Args:
        x: Input tensor (B, T, D)
        gamma: Scale parameter (D,)
        beta: Bias parameter (D,)
        dy: Upstream gradient (B, T, D)
        h: Step size

    Returns:
        Numerical gradient for gamma (D,)
    """
    D = len(gamma)
    num_grad = np.zeros(D)

    for i in range(D):
        # Save original value
        original = gamma[i]

        # f(gamma + h)
        gamma[i] = original + h
        y_plus, _ = layer_norm_forward(x, gamma, beta)
        loss_plus = np.sum(y_plus * dy)

        # f(gamma - h)
        gamma[i] = original - h
        y_minus, _ = layer_norm_forward(x, gamma, beta)
        loss_minus = np.sum(y_minus * dy)

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

        # Restore original
        gamma[i] = original

    return num_grad

def compute_numerical_gradient_beta(
    x: np.ndarray,
    gamma: np.ndarray,
    beta: np.ndarray,
    dy: np.ndarray,
    h: float = 1e-5
) -> np.ndarray:
    """
    Compute numerical gradient for beta using finite differences.

    Args:
        x: Input tensor (B, T, D)
        gamma: Scale parameter (D,)
        beta: Bias parameter (D,)
        dy: Upstream gradient (B, T, D)
        h: Step size

    Returns:
        Numerical gradient for beta (D,)
    """
    D = len(beta)
    num_grad = np.zeros(D)

    for i in range(D):
        # Save original value
        original = beta[i]

        # f(beta + h)
        beta[i] = original + h
        y_plus, _ = layer_norm_forward(x, gamma, beta)
        loss_plus = np.sum(y_plus * dy)

        # f(beta - h)
        beta[i] = original - h
        y_minus, _ = layer_norm_forward(x, gamma, beta)
        loss_minus = np.sum(y_minus * dy)

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

        # Restore original
        beta[i] = original

    return num_grad

def compute_numerical_gradient_x(
    x: np.ndarray,
    gamma: np.ndarray,
    beta: np.ndarray,
    dy: np.ndarray,
    h: float = 1e-5,
    max_elements: int = 100000
) -> np.ndarray:
    """
    Compute numerical gradient for x using finite differences.
    For large tensors, uses spot check.

    Args:
        x: Input tensor (B, T, D)
        gamma: Scale parameter (D,)
        beta: Bias parameter (D,)
        dy: Upstream gradient (B, T, D)
        h: Step size
        max_elements: Maximum elements to check (spot check if larger)

    Returns:
        Numerical gradient for x (B, T, D) or spot check values
    """
    B, T, D = x.shape
    total_elements = B * T * D

    if total_elements <= max_elements:
        # Full gradient check
        num_grad = np.zeros_like(x)
        x_flat = x.flatten()
        num_grad_flat = num_grad.flatten()

        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, _ = layer_norm_forward(x, gamma, beta)
            loss_plus = np.sum(y_plus * dy)

            # f(x - h)
            x_flat[i] = original - h
            y_minus, _ = layer_norm_forward(x, gamma, beta)
            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
    else:
        # Spot check
        print(f"  Spot checking {max_elements} random elements...")
        n_samples = max_elements
        indices = [tuple(np.random.randint(b) for b in (B, T, D)) for _ in range(n_samples)]

        num_grad = np.zeros_like(x)
        for idx in indices:
            bi, ti, di = idx
            original = x[bi, ti, di]

            # f(x + h)
            x[bi, ti, di] = original + h
            y_plus, _ = layer_norm_forward(x, gamma, beta)
            loss_plus = np.sum(y_plus * dy)

            # f(x - h)
            x[bi, ti, di] = original - h
            y_minus, _ = layer_norm_forward(x, gamma, beta)
            loss_minus = np.sum(y_minus * dy)

            num_grad[bi, ti, di] = (loss_plus - loss_minus) / (2 * h)

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

        return num_grad

def gradient_check(
    x: np.ndarray,
    gamma: np.ndarray,
    beta: 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|

    Args:
        x: Input tensor
        gamma: Scale parameter
        beta: Bias parameter
        dy: Upstream gradient
        rtol: Relative tolerance
        atol: Absolute tolerance
        verbose: Print detailed results

    Returns:
        Dictionary of pass/fail for each parameter
    """
    results = {}

    # Store originals
    orig_gamma = gamma.copy()
    orig_beta = beta.copy()
    orig_x = x.copy()

    # -------------------------------------------------------------------------
    # Forward pass to get analytical gradients
    # -------------------------------------------------------------------------
    y, cache = layer_norm_forward(x, gamma, beta)
    dx_analytical, d_gamma_analytical, d_beta_analytical = layer_norm_backward(dy, cache)

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

    # Reset gamma to original
    gamma[:] = orig_gamma

    d_gamma_numerical = compute_numerical_gradient_gamma(x, gamma, beta, dy)

    # Compare
    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)

    # Reset beta to original
    beta[:] = orig_beta

    d_beta_numerical = compute_numerical_gradient_beta(x, gamma, beta, dy)

    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)

    # Reset x to original
    x[:] = orig_x

    d_x_numerical = compute_numerical_gradient_x(x, gamma, beta, 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

    # Restore originals
    gamma[:] = orig_gamma
    beta[:] = orig_beta
    x[:] = orig_x

    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 inputs
        np.random.seed(123)
        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)

        # Run gradient check
        results = gradient_check(x, gamma, beta, 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 39559 bytes to ./minimax-m2.7/backwards/layer_norm_numpy.py


Now let me run the updated code to test the gradient checks.

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
============================================================
Analytical gradient shape: (8,)
Numerical gradient shape: (8,)
Max absolute difference: 8.55e-11
Max relative tolerance: 4.55e-04
Mean analytical: 2.429868e+00
Mean numerical: 2.429868e+00

Gradient check: PASSED ✓

============================================================
GRADIENT CHECK: beta
============================================================
Analytical gradient shape: (8,)
Numerical gradient shape: (8,)
Max absolute difference: 1.08e-10
Max relative tolerance: 5.13e-04
Mean analytical: 2.602642e+00
Mean numerical: 2.602642e+00

Gradient check: PASSED ✓

============================================================
GRADIENT CHECK: x (input)
============================================================
Analytical gradient shape: (2, 4, 8)
Numerical gradient shape: (2, 4, 8)
Max absolute difference: 2.14e+00
Max relative tolerance: 1.00e-05
Mean analytical: 5.202041e-01
Mean numerical: 0.000000e+00

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

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

============================================================
GRADIENT CHECK: gamma
============================================================
Analytical gradient shape: (16,)
Numerical gradient shape: (16,)
Max absolute difference: 5.52e-10
Max relative tolerance: 1.28e-03
Mean analytical: 4.142521e+00
Mean numerical: 4.142521e+00

Gradient check: PASSED ✓

============================================================
GRADIENT CHECK: beta
============================================================
Analytical gradient shape: (16,)
Numerical gradient shape: (16,)
Max absolute difference: 2.52e-10
Max relative tolerance: 8.59e-04
Mean analytical: 3.506940e+00
Mean numerical: 3.506940e+00

Gradient check: PASSED ✓

============================================================
GRADIENT CHECK: x (input)
============================================================
Analytical gradient shape: (4, 8, 16)
Numerical gradient shape: (4, 8, 16)
Max absolute difference: 4.60e+00
Max relative tolerance: 1.00e-05
Mean analytical: 6.141263e-01
Mean numerical: 0.000000e+00

Gradient check: FAILED ✗
     gamma: ✓ PASS
      beta: ✓ PASS
         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.38e-09
Max relative tolerance: 2.84e-03
Mean analytical: 9.330259e+00
Mean numerical: 9.330259e+00

Gradient check: PASSED ✓

============================================================
GRADIENT CHECK: beta
============================================================
Analytical gradient shape: (32,)
Numerical gradient shape: (32,)
Max absolute difference: 2.00e-09
Max relative tolerance: 2.61e-03
Mean analytical: 9.296132e+00
Mean numerical: 9.296132e+00

Gradient check: PASSED ✓

============================================================
GRADIENT CHECK: x (input)
============================================================
Analytical gradient shape: (8, 16, 32)
Numerical gradient shape: (8, 16, 32)
Max absolute difference: 6.83e+00
Max relative tolerance: 1.00e-05
Mean analytical: 7.088609e-01
Mean numerical: 0.000000e+00

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

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

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

Shape: (B=32, T=128, D=256)
----------------------------------------
Forward:  6.505 ± 0.966 ms (0.2 GB/s)
Backward: 8.042 ± 0.830 ms (0.1 GB/s)
Total:    14.547 ms

Shape: (B=64, T=128, D=512)
----------------------------------------
Forward:  29.728 ± 3.325 ms (0.1 GB/s)
Backward: 38.809 ± 4.349 ms (0.1 GB/s)
Total:    68.537 ms

Shape: (B=32, T=512, D=768)
----------------------------------------
Forward:  84.739 ± 8.603 ms (0.1 GB/s)
Backward: 120.727 ± 9.439 ms (0.1 GB/s)
Total:    205.466 ms

Shape: (B=16, T=512, D=1024)
----------------------------------------
Forward:  60.770 ± 4.044 ms (0.1 GB/s)
Backward: 80.685 ± 7.889 ms (0.1 GB/s)
Total:    141.455 ms


The gradient check for x is still failing. The numerical gradient is all zeros, but the analytical gradient is non-zero. This is suspicious. Let me check the numerical gradient computation for x more carefully.

Looking at the output: `Mean numerical: 0.000000e+00` - the numerical gradient is zero!

But the mean analytical is non-zero: `Mean analytical: 5.202041e-01`

This suggests something is wrong with the numerical gradient computation. Let me trace through what happens when we modify x.

In `compute_numerical_gradient_x`, we modify x in-place:
```python
x_flat[i] = original + h
y_plus, _ = layer_norm_forward(x, gamma, beta)
```

But wait, in `layer_norm_forward`, we compute:
```python
x_centered = x - mean  # (B, T, D)
```

Where `mean = np.mean(x, axis=-1, keepdims=True)`. This is computed over the original x values.

But if we modify x in-place, we're changing it. So the next iteration starts with the modified x.

The issue is that we're modifying the array and not properly restoring it between iterations!

Actually, looking more carefully:
```python
# Restore
x_flat[i] = original
```

This should restore it. Let me trace through more carefully...

Actually wait - when we modify x for one element, the mean changes! So the next element's numerical gradient computation is affected by the previous modification.

But that shouldn't result in all zeros... Let me trace through a single iteration:

1. Get original value at index i
2. Set x_flat[i] = original + h
3. Forward pass: y_plus, _ = layer_norm_forward(x, gamma, beta)
   - mean = np.mean(x, axis=-1, keepdims=True)
   - x_centered = x - mean
   - etc.
4. loss_plus = np.sum(y_plus * dy)
5. Set x_flat[i] = original - h
6. Forward pass: y_minus, _ = layer_norm_forward(x, gamma, beta)
7. loss_minus = np.sum(y_minus * dy)
8. num_grad_flat[i] = (loss_plus - loss_minus) / (2 * h)
9. x_flat[i] = original (RESTORE)

This should work correctly for a single element. Let me add some debugging to verify the numerical gradient computation is actually running and producing non-zero results.

Actually, I think the problem might be in how the finite difference approximation handles the x gradient. When I compute the loss, it's the sum of y * dy. For a tiny perturbation h = 1e-5, the change in loss should be approximately the gradient. But if the loss difference is negligible relative to h, that could explain the zero numerical gradient.

Wait—I think I found it. In `gradient_check`, I'm copying x with `orig_x = x.copy()`, then calling `compute_numerical_gradient_x` which modifies x in place. The restoration at the end with `x[:] = orig_x` should handle this, but the numerical gradient itself might be computed before any modifications persist.

Let me trace through: after the gamma check, gamma gets reset to original values. Then for the x check, x is reset to original, so the next call to `compute_numerical_gradient_x` starts from the correct state. But I need to verify the loss function is actually producing non-zero differences when I perturb x.

Actually, the real problem might be that the numerical gradient computation is working correctly but the loss changes are so small relative to machine precision that they round to zero, while the analytical gradient has non-trivial values.

Let me trace through what happens when I perturb x[0,0,0] - the forward pass recomputes the mean and variance, which changes the normalized output y. The loss should change, but if the perturbation is too small relative to the data scale, the difference could fall below the numerical precision threshold. The mean gradient for dy at that position might be small enough that the effect is negligible, causing the numerical gradient to compute as zero. Let me verify this with a direct check.

Let me debug the numerical gradient computation for x:

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))

    gamma = np.ones(D)  # assume identity gamma for simplicity
    dz = dy * gamma
    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)
beta = np.zeros(D).astype(np.float64)
dy = np.random.randn(B, T, D).astype(np.float64)

# Forward
y, cache = layer_norm_forward(x, gamma, beta)
dx_analytical, d_gamma_analytical, d_beta_analytical = layer_norm_backward(dy, cache)

print('Analytical dx[0,0,0]:', dx_analytical[0, 0, 0])

# Numerical gradient for x[0,0,0] only
h = 1e-5

def loss_fn(x_input):
    y_out, _ = layer_norm_forward(x_input, gamma, beta)
    return np.sum(y_out * dy)

loss_base = loss_fn(x)
print('Base loss:', loss_base)

x_plus = x.copy()
x_plus[0, 0, 0] += h
loss_plus = loss_fn(x_plus)
print('Loss with x[0,0,0] + h:', loss_plus)

x_minus = x.copy()
x_minus[0, 0, 0] -= h
loss_minus = loss_fn(x_minus)
print('Loss with x[0,0,0] - h:', loss_minus)

d_x_num = (loss_plus - loss_minus) / (2 * h)
print('Numerical dx[0,0,0]:', d_x_num)

# But wait - when x changes, mean changes too!
# Let me check if that's the issue
print()
print('Checking mean dependency...')
print('Original mean[0,0,0]:', cache['mean'][0,0,0])

# Compute numerical gradient properly
# When x changes, mean changes too
print()
print('Key insight: when x[0,0,0] changes, the mean over D for position (0,0) changes')
print('This is captured correctly in the forward pass when we recompute')
"


Analytical dx[0,0,0]: 0.3440944762646291
Base loss: -3.607272590372364
Loss with x[0,0,0] + h: -3.607269149442598
Loss with x[0,0,0] - h: -3.607276031332122
Numerical dx[0,0,0]: 0.34409447620120653

Checking mean dependency...
Original mean[0,0,0]: 0.5509406824170687

Key insight: when x[0,0,0] changes, the mean over D for position (0,0) changes
This is captured correctly in the forward pass when we recompute


Interesting - when I run this correctly (resetting x to a new array, not modifying in-place), I get matching results. The issue in my gradient check must be with how I'm handling the in-place modification.

Let me debug the in-place version to see what's going wrong.

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))

    gamma = np.ones(D)  # assume identity gamma for simplicity
    dz = dy * gamma
    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

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)
beta = np.zeros(D).astype(np.float64)
dy = np.random.randn(B, T, D).astype(np.float64)

# Test in-place modification
x_test = x.copy()
print('Original x[0,0,0]:', x_test[0,0,0])

h = 1e-5

# f(x + h)
original = x_test[0,0,0]
x_test[0,0,0] = original + h
y_plus, _ = layer_norm_forward(x_test, gamma, beta)
loss_plus = np.sum(y_plus * dy)
print('Loss plus:', loss_plus)

# f(x - h)
x_test[0,0,0] = original - h
y_minus, _ = layer_norm_forward(x_test, gamma, beta)
loss_minus = np.sum(y_minus * dy)
print('Loss minus:', loss_minus)

# Restore
x_test[0,0,0] = original

d_x_num = (loss_plus - loss_minus) / (2 * h)
print('Numerical dx[0,0,0]:', d_x_num)

# Now compare with analytical
y, cache = layer_norm_forward(x, gamma, beta)
dx_analytical, _, _ = layer_norm_backward(dy, cache)
print('Analytical dx[0,0,0]:', dx_analytical[0,0,0])
print('Difference:', abs(d_x_num - dx_analytical[0,0,0]))
"


Original x[0,0,0]: 0.4967141530112327
Loss plus: -3.607269149442598
Loss minus: -3.607276031332122
Numerical dx[0,0,0]: 0.34409447620120653
Analytical dx[0,0,0]: 0.3440944762646291
Difference: 6.342254499358546e-11


So the gradient check passes when I use the correct setup. The issue in the main code must be related to how we're handling the in-place modifications with x. Let me look at the `compute_numerical_gradient_x` function more carefully.

Actually, the issue is clear now - in `compute_numerical_gradient_x`, we do `x_flat[i] = original + h` which modifies x in place, but x is passed by reference. Then on the next iteration, when we try to modify another element, the mean has already changed from the previous perturbation!

Let me fix this by always using x.copy() or restoring from the original each time.


    return num_grad

def compute_numerical_gradient_beta(
    x: np.ndarray,
    gamma: np.ndarray,
    beta: np.ndarray,
    dy: np.ndarray,
    h: float = 1e-5
) -> np.ndarray:
    """
    Compute numerical gradient for beta using finite differences.

    Args:
        x: Input tensor (B, T, D)
        gamma: Scale parameter (D,)
        beta: Bias parameter (D,)
        dy: Upstream gradient (B, T, D)
        h: Step size

    Returns:
        Numerical gradient for beta (D,)
    """
    D = len(beta)
    num_grad = np.zeros(D)

    for i in range(D):
        # Save original value
        original = beta[i]

        # f(beta + h)
        beta[i] = original + h
        y_plus, _ = layer_norm_forward(x, gamma, beta)
        loss_plus = np.sum(y_plus * dy)

        # f(beta - h)
        beta[i] = original - h
        y_minus, _ = layer_norm_forward(x, gamma, beta)
        loss_minus = np.sum(y_minus * dy)

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

        # Restore original
        beta[i] = original

    return num_grad

def compute_numerical_gradient_x(
    x: np.ndarray,
    gamma: np.ndarray,
    beta: np.ndarray,
    dy: np.ndarray,
    h: float = 1e-5,
    max_elements: int = 100000
) -> np.ndarray:
    """
    Compute numerical gradient for x using finite differences.
    For large tensors, uses spot check.

    Args:
        x: Input tensor (B, T, D)
        gamma: Scale parameter (D,)
        beta: Bias parameter (D,)
        dy: Upstream gradient (B, T, D)
        h: Step size
        max_elements: Maximum elements to check (spot check if larger)

    Returns:

[534 more lines in file. Use offset=410 to continue.]


I see the issue now. The numerical gradient functions modify x/gamma/beta in place but don't restore them properly to their original values at the end. The last test modifies x but doesn't restore, then when the analytical gradient is computed, x is in a modified state. Let me fix this by ensuring proper restoration at the end of each function.

edit

def compute_numerical_gradient_x(
    x: np.ndarray,
    gamma: np.ndarray,
    beta: np.ndarray,
    dy: np.ndarray,
    h: float = 1e-5,
    max_elements: int = 100000
) -> Tuple[np.ndarray, np.ndarray]:
    """
    Compute numerical gradient for x using finite differences.
    For large tensors, uses spot check.

    Returns both the numerical gradient AND the numerical gradient for spot-checked elements.

    Args:
        x: Input tensor (B, T, D) - will be restored to original values
        gamma: Scale parameter (D,)
        beta: Bias parameter (D,)
        dy: Upstream gradient (B, T, D)
        h: Step size
        max_elements: Maximum elements to check (spot check if larger)

    Returns:
        Tuple of (num_grad, spot_check_mask) where spot_check_mask marks checked elements
    """
    B, T, D = x.shape
    total_elements = B * T * D

    # Save original x values
    orig_x = x.copy()

    if total_elements <= max_elements:
        # Full gradient check
        num_grad = np.zeros_like(x)
        x_flat = x.flatten()
        num_grad_flat = num_grad.flatten()

        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, _ = layer_norm_forward(x, gamma, beta)
            loss_plus = np.sum(y_plus * dy)

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

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

            # Restore
            x_flat[i] = original

        # Restore x to original values
        x[:] = orig_x

        return num_grad, np.ones((B, T, D), dtype=bool)  # All elements checked
    else:
        # Spot check
        print(f"  Spot checking {max_elements} random elements...")
        n_samples = max_elements
        num_grad = np.zeros_like(x)
        spot_checked = np.zeros((B, T, D), dtype=bool)

        indices = [tuple(np.random.randint(b) for b in (B, T, D)) for _ in range(n_samples)]

        for idx in indices:
            bi, ti, di = idx
            original = x[bi, ti, di]
            spot_checked[bi, ti, di] = True

            # f(x + h)
            x[bi, ti, di] = original + h
            y_plus, _ = layer_norm_forward(x, gamma, beta)
            loss_plus = np.sum(y_plus * dy)

            # f(x - h)
            x[bi, ti, di] = original - h
            y_minus, _ = layer_norm_forward(x, gamma, beta)
            loss_minus = np.sum(y_minus * dy)

            num_grad[bi, ti, di] = (loss_plus - loss_minus) / (2 * h)

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

        # Restore x to original values
        x[:] = orig_x

        return num_grad, spot_checked

def compute_numerical_gradient_x(
    x: np.ndarray,
    gamma: np.ndarray,
    beta: np.ndarray,
    dy: np.ndarray,
    h: float = 1e-5,
    max_elements: int = 100000
) -> np.ndarray:
    """
    Compute numerical gradient for x using finite differences.
    For large tensors, uses spot check.

    Args:
        x: Input tensor (B, T, D)
        gamma: Scale parameter (D,)
        beta: Bias parameter (D,)
        dy: Upstream gradient (B, T, D)
        h: Step size
        max_elements: Maximum elements to check (spot check if larger)

    Returns:
        Numerical gradient for x (B, T, D) or spot check values
    """
    B, T, D = x.shape
    total_elements = B * T * D

    if total_elements <= max_elements:
        # Full gradient check
        num_grad = np.zeros_like(x)
        x_flat = x.flatten()
        num_grad_flat = num_grad.flatten()

        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, _ = layer_norm_forward(x, gamma, beta)
            loss_plus = np.sum(y_plus * dy)

            # f(x - h)
            x_flat[i] = original - h
            y_minus, _ = layer_norm_forward(x, gamma, beta)
            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
    else:
        # Spot check
        print(f"  Spot checking {max_elements} random elements...")
        n_samples = max_elements
        indices = [tuple(np.random.randint(b) for b in (B, T, D)) for _ in range(n_samples)]

        num_grad = np.zeros_like(x)
        for idx in indices:
            bi, ti, di = idx
            original = x[bi, ti, di]

            # f(x + h)
            x[bi, ti, di] = original + h
            y_plus, _ = layer_norm_forward(x, gamma, beta)
            loss_plus = np.sum(y_plus * dy)

            # f(x - h)
            x[bi, ti, di] = original - h
            y_minus, _ = layer_norm_forward(x, gamma, beta)
            loss_minus = np.sum(y_minus * dy)

            num_grad[bi, ti, di] = (loss_plus - loss_minus) / (2 * h)

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

        return num_grad


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

     ...
 392     beta: np.ndarray,
 393     dy: np.ndarray,
 394     h: float = 1e-5,
 395     max_elements: int = 100000
-396 ) -> np.ndarray:
+396 ) -> Tuple[np.ndarray, np.ndarray]:
 397     """
 398     Compute numerical gradient for x using finite differences.
 399     For large tensors, uses spot check.
 400
+401     Returns both the numerical gradient AND the numerical gradient for spot-checked elements.
+402
 401     Args:
-402         x: Input tensor (B, T, D)
+404         x: Input tensor (B, T, D) - will be restored to original values
 403         gamma: Scale parameter (D,)
 404         beta: Bias parameter (D,)
 405         dy: Upstream gradient (B, T, D)
 406         h: Step size
 407         max_elements: Maximum elements to check (spot check if larger)
 408
 409     Returns:
-410         Numerical gradient for x (B, T, D) or spot check values
+412         Tuple of (num_grad, spot_check_mask) where spot_check_mask marks checked elements
 411     """
 412     B, T, D = x.shape
 413     total_elements = B * T * D
 414
+417     # Save original x values
+418     orig_x = x.copy()
+419
 415     if total_elements <= max_elements:
 416         # Full gradient check
 417         num_grad = np.zeros_like(x)
 418         x_flat = x.flatten()
     ...
 439
 440             # Restore
 441             x_flat[i] = original
 442
-443         return num_grad
+448         # Restore x to original values
+449         x[:] = orig_x
+450
+451         return num_grad, np.ones((B, T, D), dtype=bool)  # All elements checked
 444     else:
 445         # Spot check
 446         print(f"  Spot checking {max_elements} random elements...")
 447         n_samples = max_elements
+456         num_grad = np.zeros_like(x)
+457         spot_checked = np.zeros((B, T, D), dtype=bool)
+458
 448         indices = [tuple(np.random.randint(b) for b in (B, T, D)) for _ in range(n_samples)]
 449
-450         num_grad = np.zeros_like(x)
 451         for idx in indices:
 452             bi, ti, di = idx
 453             original = x[bi, ti, di]
+464             spot_checked[bi, ti, di] = True
 454
 455             # f(x + h)
 456             x[bi, ti, di] = original + h
 457             y_plus, _ = layer_norm_forward(x, gamma, beta)
     ...
 466
 467             # Restore
 468             x[bi, ti, di] = original
 469
-470         return num_grad
+481         # Restore x to original values
+482         x[:] = orig_x
+483
+484         return num_grad, spot_checked
 471
 472
 473 def gradient_check(
 474     x: np.ndarray,
     ...


Now I need to update the gradient_check function to use the new signature and compare only spot-checked elements.


    Args:
        x: Input tensor
        gamma: Scale parameter
        beta: Bias parameter
        dy: Upstream gradient
        rtol: Relative tolerance
        atol: Absolute tolerance
        verbose: Print detailed results

    Returns:
        Dictionary of pass/fail for each parameter
    """
    results = {}

    # Store originals
    orig_gamma = gamma.copy()
    orig_beta = beta.copy()
    orig_x = x.copy()

    # -------------------------------------------------------------------------
    # Forward pass to get analytical gradients
    # -------------------------------------------------------------------------
    y, cache = layer_norm_forward(x, gamma, beta)
    dx_analytical, d_gamma_analytical, d_beta_analytical = layer_norm_backward(dy, cache)

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

    # Reset gamma to original
    gamma[:] = orig_gamma

    d_gamma_numerical = compute_numerical_gradient_gamma(x, gamma, beta, dy)

    # Compare
    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}")

[408 more lines in file. Use offset=550 to continue.]


Now I need to update the gradient_check function to use the new return type of compute_numerical_gradient_x. The gradient check for x needs to compare only the elements that were actually checked (those that have numerical gradients).


    results['beta'] = passed

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

    # Reset x to original
    x[:] = orig_x

    d_x_numerical = compute_numerical_gradient_x(x, gamma, beta, 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

    # Restore originals
    gamma[:] = orig_gamma
    beta[:] = orig_beta
    x[:] = orig_x

    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)            │ -                             ║

[318 more lines in file. Use offset=640 to continue.]


Now I need to update the gradient_check function to handle the new return type of compute_numerical_gradient_x. The function now returns a tuple (num_grad, spot_checked_mask), so I need to update the code to handle this.

    # Reset x to original
    x[:] = orig_x

    d_x_numerical, spot_checked = compute_numerical_gradient_x(x, gamma, beta, dy)

    # Only check elements that were numerically computed
    diff = np.abs(dx_analytical[spot_checked] - d_x_numerical[spot_checked])
    tolerance = atol + rtol * np.abs(d_x_numerical[spot_checked])
    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"Elements checked: {np.sum(spot_checked)} / {spot_checked.size}")
        if np.any(spot_checked):
            print(f"Max absolute difference: {np.max(diff):.2e}")
            print(f"Max relative tolerance: {np.max(tolerance):.2e}")
            print(f"Mean analytical (checked): {np.mean(np.abs(dx_analytical[spot_checked])):.6e}")
            print(f"Mean numerical (checked): {np.mean(np.abs(d_x_numerical[spot_checked])):.6e}")
        print(f"\nGradient check: {'PASSED ✓' if passed else 'FAILED ✗'}")

    # Reset x to original
    x[:] = orig_x

    d_x_numerical = compute_numerical_gradient_x(x, gamma, beta, 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 ✗'}")


     ...
 590
 591     # Reset x to original
 592     x[:] = orig_x
 593
-594     d_x_numerical = compute_numerical_gradient_x(x, gamma, beta, dy)
+594     d_x_numerical, spot_checked = compute_numerical_gradient_x(x, gamma, beta, dy)
 595
-596     diff = np.abs(dx_analytical - d_x_numerical)
-597     tolerance = atol + rtol * np.abs(d_x_numerical)
+596     # Only check elements that were numerically computed
+597     diff = np.abs(dx_analytical[spot_checked] - d_x_numerical[spot_checked])
+598     tolerance = atol + rtol * np.abs(d_x_numerical[spot_checked])
 598     passed = np.all(diff <= tolerance)
 599
 600     if verbose:
 601         print(f"Analytical gradient shape: {dx_analytical.shape}")
 602         print(f"Numerical gradient shape: {d_x_numerical.shape}")
-603         print(f"Max absolute difference: {np.max(diff):.2e}")
-604         print(f"Max relative tolerance: {np.max(tolerance):.2e}")
-605         print(f"Mean analytical: {np.mean(np.abs(dx_analytical)):.6e}")
-606         print(f"Mean numerical: {np.mean(np.abs(d_x_numerical)):.6e}")
+604         print(f"Elements checked: {np.sum(spot_checked)} / {spot_checked.size}")
+605         if np.any(spot_checked):
+606             print(f"Max absolute difference: {np.max(diff):.2e}")
+607             print(f"Max relative tolerance: {np.max(tolerance):.2e}")
+608             print(f"Mean analytical (checked): {np.mean(np.abs(dx_analytical[spot_checked])):.6e}")
+609             print(f"Mean numerical (checked): {np.mean(np.abs(d_x_numerical[spot_checked])):.6e}")
 607         print(f"\nGradient check: {'PASSED ✓' if passed else 'FAILED ✗'}")
 608
 609     results['x'] = passed
 610
     ...


Good. Now let me also add the Tuple import since I used it in the type hint:

from typing import Tuple, Dict, Optional


No changes made to ./minimax-m2.7/backwards/layer_norm_numpy.py. The replacement produced identical content. This might indicate an issue with special characters or the text not existing as expected.


The import is already there. Let me run the code to see if it works now.

bash

cd ./minimax-m2.7/backwards && python layer_norm_numpy.py 2>&1 | head -150