Compare reference and optimized outputs. Returns comparison metrics.
(
ref_output: torch.Tensor,
opt_output: torch.Tensor,
dtype: torch.dtype,
custom_atol: Optional[float] = None,
custom_rtol: Optional[float] = None,
)
| 685 | |
| 686 | |
| 687 | def compare_outputs( |
| 688 | ref_output: torch.Tensor, |
| 689 | opt_output: torch.Tensor, |
| 690 | dtype: torch.dtype, |
| 691 | custom_atol: Optional[float] = None, |
| 692 | custom_rtol: Optional[float] = None, |
| 693 | ) -> Dict[str, Any]: |
| 694 | """ |
| 695 | Compare reference and optimized outputs. Returns comparison metrics. |
| 696 | """ |
| 697 | result: Dict[str, Any] = {} |
| 698 | |
| 699 | # Shape check |
| 700 | result["shapes_match"] = ref_output.shape == opt_output.shape |
| 701 | result["ref_shape"] = str(list(ref_output.shape)) |
| 702 | result["opt_shape"] = str(list(opt_output.shape)) |
| 703 | |
| 704 | if not result["shapes_match"]: |
| 705 | result["correctness"] = "FAIL" |
| 706 | result["reason"] = f"Shape mismatch: ref={result['ref_shape']}, opt={result['opt_shape']}" |
| 707 | return result |
| 708 | |
| 709 | # NaN / Inf check |
| 710 | ref_float = ref_output.float() |
| 711 | opt_float = opt_output.float() |
| 712 | |
| 713 | result["ref_has_nan"] = bool(torch.isnan(ref_float).any()) |
| 714 | result["ref_has_inf"] = bool(torch.isinf(ref_float).any()) |
| 715 | result["opt_has_nan"] = bool(torch.isnan(opt_float).any()) |
| 716 | result["opt_has_inf"] = bool(torch.isinf(opt_float).any()) |
| 717 | |
| 718 | if result["opt_has_nan"] and not result["ref_has_nan"]: |
| 719 | result["correctness"] = "FAIL" |
| 720 | result["reason"] = "Optimized output contains NaN where reference does not" |
| 721 | return result |
| 722 | |
| 723 | if result["opt_has_inf"] and not result["ref_has_inf"]: |
| 724 | result["correctness"] = "FAIL" |
| 725 | result["reason"] = "Optimized output contains Inf where reference does not" |
| 726 | return result |
| 727 | |
| 728 | # Numerical comparison |
| 729 | diff = (ref_float - opt_float).abs() |
| 730 | |
| 731 | # Mask out positions where both are NaN (those are fine) |
| 732 | valid_mask = ~(torch.isnan(ref_float) & torch.isnan(opt_float)) |
| 733 | if valid_mask.any(): |
| 734 | valid_diff = diff[valid_mask] |
| 735 | result["max_abs_error"] = float(valid_diff.max()) |
| 736 | result["mean_abs_error"] = float(valid_diff.mean()) |
| 737 | else: |
| 738 | result["max_abs_error"] = 0.0 |
| 739 | result["mean_abs_error"] = 0.0 |
| 740 | |
| 741 | # Tolerance check |
| 742 | tols = DEFAULT_TOLERANCES.get(dtype, {"atol": 1e-4, "rtol": 1e-4}) |
| 743 | atol = custom_atol if custom_atol is not None else tols["atol"] |
| 744 | rtol = custom_rtol if custom_rtol is not None else tols["rtol"] |
no outgoing calls
no test coverage detected