Spaces:
Running on Zero
Running on Zero
| 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 | |