Compare two tensors. Returns match info.
(output, expected, atol: float, rtol: float)
| 166 | |
| 167 | |
| 168 | def _compare(output, expected, atol: float, rtol: float) -> Dict[str, Any]: |
| 169 | """Compare two tensors. Returns match info.""" |
| 170 | import torch |
| 171 | |
| 172 | result: Dict[str, Any] = { |
| 173 | "match": False, |
| 174 | "reason": "", |
| 175 | "max_abs_error": float("inf"), |
| 176 | "mean_abs_error": float("inf"), |
| 177 | } |
| 178 | |
| 179 | # Shape check |
| 180 | if output.shape != expected.shape: |
| 181 | result["reason"] = f"shape mismatch: {output.shape} vs {expected.shape}" |
| 182 | return result |
| 183 | |
| 184 | # Cast to float32 for comparison |
| 185 | out_f = output.detach().float().cpu() |
| 186 | exp_f = expected.detach().float().cpu() |
| 187 | |
| 188 | # NaN/Inf symmetry |
| 189 | out_nan = torch.isnan(out_f) |
| 190 | exp_nan = torch.isnan(exp_f) |
| 191 | if out_nan.any() or exp_nan.any(): |
| 192 | if not torch.equal(out_nan, exp_nan): |
| 193 | result["reason"] = "NaN position mismatch" |
| 194 | return result |
| 195 | mask = ~out_nan |
| 196 | if mask.any(): |
| 197 | out_f = out_f[mask] |
| 198 | exp_f = exp_f[mask] |
| 199 | else: |
| 200 | result["match"] = True |
| 201 | result["reason"] = "all NaN (matching)" |
| 202 | result["max_abs_error"] = 0.0 |
| 203 | result["mean_abs_error"] = 0.0 |
| 204 | return result |
| 205 | |
| 206 | abs_err = (out_f - exp_f).abs() |
| 207 | result["max_abs_error"] = float(abs_err.max()) |
| 208 | result["mean_abs_error"] = float(abs_err.mean()) |
| 209 | |
| 210 | if torch.allclose(out_f, exp_f, atol=atol, rtol=rtol): |
| 211 | result["match"] = True |
| 212 | result["reason"] = "PASS" |
| 213 | else: |
| 214 | result["reason"] = ( |
| 215 | f"tolerance exceeded: max_abs={result['max_abs_error']:.6e}, " |
| 216 | f"mean_abs={result['mean_abs_error']:.6e} (atol={atol}, rtol={rtol})" |
| 217 | ) |
| 218 | |
| 219 | return result |
| 220 | |
| 221 | |
| 222 | def _compare_outputs(output, expected, atol: float, rtol: float) -> Dict[str, Any]: |