Function summarize_kernel_drift

rl_engine/testing/reference_ops.py:109–142

Summarize candidate-vs-reference drift for benchmark/test output.

(
    candidate: torch.Tensor,
    reference: torch.Tensor,
    mask: torch.Tensor | None = None,
)

def summarize_kernel_drift(
    candidate: torch.Tensor,
    reference: torch.Tensor,
    mask: torch.Tensor | None = None,
) -> dict[str, Any]:
    """Summarize candidate-vs-reference drift for benchmark/test output."""

    if candidate.shape != reference.shape:
        raise ValueError(
            f"candidate shape {tuple(candidate.shape)} must match reference shape "
            f"{tuple(reference.shape)}"
        )

    diff = (candidate.float() - reference.float()).abs()
    if mask is not None:
        active = _bool_mask(mask, device=diff.device)
        active_diff = diff[active]
        active_count = int(active.sum().item())
    else:
        active_diff = diff.reshape(-1)
        active_count = int(diff.numel())

    if active_count == 0:
        max_abs = 0.0
        mean_abs = 0.0
    else:
        max_abs = float(active_diff.max().item())
        mean_abs = float(active_diff.mean().item())

    return {
        "max_abs_error": max_abs,
        "mean_abs_error": mean_abs,
        "active_count": active_count,
    }