from __future__ import annotations import math from typing import Any, Tuple, Optional try: import torch HAS_TORCH = True except ImportError: HAS_TORCH = False try: import numpy as np HAS_NUMPY = True except ImportError: HAS_NUMPY = False def summarize_obj(obj: Any, max_len: int = 60) -> str: """Generate a clean, informative one-line summary of a tensor, array, or object.""" if HAS_TORCH and isinstance(obj, torch.Tensor): return f"Tensor(shape={list(obj.shape)}, dtype={str(obj.dtype).replace('torch.', '')})" if HAS_NUMPY and isinstance(obj, np.ndarray): return f"ndarray(shape={list(obj.shape)}, dtype={obj.dtype})" if isinstance(obj, (list, tuple)): items_preview = ", ".join(summarize_obj(x, 20) for x in obj[:3]) if len(obj) > 3: items_preview += f", ... (+{len(obj)-3} more)" bracket = ("(", ")") if isinstance(obj, tuple) else ("[", "]") return f"{bracket[0]}{items_preview}{bracket[1]}" s = str(obj) if len(s) > max_len: return s[:max_len] + "..." return s def compare_outputs( actual: Any, expected: Any, atol: float = 1e-4, rtol: float = 1e-4, ) -> Tuple[bool, Optional[str], Optional[float]]: """ Compare actual output against expected output with numerical tolerance. Returns: (is_match, error_message, max_diff) """ # 1. Exact identity or both None if actual is expected: return True, None, 0.0 # 2. PyTorch Tensor comparison if HAS_TORCH and (isinstance(actual, torch.Tensor) or isinstance(expected, torch.Tensor)): if not isinstance(actual, torch.Tensor): return False, f"Type mismatch: expected torch.Tensor, got {type(actual).__name__}", None if not isinstance(expected, torch.Tensor): return False, f"Type mismatch: expected {type(expected).__name__}, got torch.Tensor", None # Check shape if actual.shape != expected.shape: return False, f"Shape mismatch: expected {list(expected.shape)}, got {list(actual.shape)}", None # Check dtype if actual.dtype != expected.dtype: return False, f"Dtype mismatch: expected {expected.dtype}, got {actual.dtype}", None act = actual.detach().cpu().float() exp = expected.detach().cpu().float() # Check for NaN / Inf if torch.isnan(act).any() and not torch.isnan(exp).any(): return False, "Output contains unexpected NaN values", None if torch.isinf(act).any() and not torch.isinf(exp).any(): return False, "Output contains unexpected Inf values", None # Filter out identical infinite values (e.g. -inf and -inf) to avoid (-inf) - (-inf) => NaN matching_neginf = torch.isneginf(act) & torch.isneginf(exp) matching_posinf = torch.isposinf(act) & torch.isposinf(exp) matching_inf = matching_neginf | matching_posinf matching_nan = torch.isnan(act) & torch.isnan(exp) ignore_mask = matching_inf | matching_nan if ignore_mask.all(): max_diff = 0.0 else: act_diff = torch.where(ignore_mask, torch.zeros_like(act), act) exp_diff = torch.where(ignore_mask, torch.zeros_like(exp), exp) diff = torch.abs(act_diff - exp_diff) finite_diff = diff[~torch.isnan(diff) & ~torch.isinf(diff)] max_diff = float(finite_diff.max().item()) if finite_diff.numel() > 0 else 0.0 if torch.allclose(act, exp, rtol=rtol, atol=atol, equal_nan=True): return True, None, max_diff else: return ( False, f"Numerical mismatch: max absolute diff {max_diff:.6e} > tolerance (atol={atol}, rtol={rtol})", max_diff, ) # 3. NumPy ndarray comparison if HAS_NUMPY and (isinstance(actual, np.ndarray) or isinstance(expected, np.ndarray)): if not isinstance(actual, np.ndarray): return False, f"Type mismatch: expected np.ndarray, got {type(actual).__name__}", None if not isinstance(expected, np.ndarray): return False, f"Type mismatch: expected {type(expected).__name__}, got np.ndarray", None if actual.shape != expected.shape: return False, f"Shape mismatch: expected {list(expected.shape)}, got {list(actual.shape)}", None if actual.dtype != expected.dtype: return False, f"Dtype mismatch: expected {expected.dtype}, got {actual.dtype}", None act = actual.astype(np.float64) exp = expected.astype(np.float64) if np.isnan(act).any() and not np.isnan(exp).any(): return False, "Output contains unexpected NaN values", None if np.isinf(act).any() and not np.isinf(exp).any(): return False, "Output contains unexpected Inf values", None matching_inf = np.isinf(act) & np.isinf(exp) & (np.sign(act) == np.sign(exp)) matching_nan = np.isnan(act) & np.isnan(exp) ignore_mask = matching_inf | matching_nan if ignore_mask.all(): max_diff = 0.0 else: act_diff = np.where(ignore_mask, 0.0, act) exp_diff = np.where(ignore_mask, 0.0, exp) diff = np.abs(act_diff - exp_diff) finite_diff = diff[~np.isnan(diff) & ~np.isinf(diff)] max_diff = float(np.max(finite_diff)) if finite_diff.size > 0 else 0.0 if np.allclose(act, exp, rtol=rtol, atol=atol, equal_nan=True): return True, None, max_diff else: return ( False, f"Numerical mismatch: max absolute diff {max_diff:.6e} > tolerance (atol={atol}, rtol={rtol})", max_diff, ) # 4. Sequences (tuple / list) if isinstance(expected, (tuple, list)): if not isinstance(actual, (tuple, list)): return False, f"Type mismatch: expected {type(expected).__name__}, got {type(actual).__name__}", None if len(actual) != len(expected): return False, f"Length mismatch: expected {len(expected)} elements, got {len(actual)}", None max_sub_diff = 0.0 for i, (act_elem, exp_elem) in enumerate(zip(actual, expected)): match, err, sub_diff = compare_outputs(act_elem, exp_elem, atol=atol, rtol=rtol) if sub_diff is not None: max_sub_diff = max(max_sub_diff, sub_diff) if not match: return False, f"Element [{i}] mismatch: {err}", max_sub_diff return True, None, max_sub_diff # 5. Scalar numbers (float, int) if isinstance(expected, (float, int)) and isinstance(actual, (float, int)): diff = abs(float(actual) - float(expected)) allowed = atol + rtol * abs(float(expected)) if diff <= allowed: return True, None, diff return False, f"Value mismatch: expected {expected}, got {actual} (diff {diff:.6e} > {allowed:.6e})", diff # 6. Fallback exact equality if actual == expected: return True, None, 0.0 return False, f"Expected {expected!r}, but got {actual!r}", None