Test numerical stability: check for NaN/Inf on normal inputs.
(
model_new,
get_inputs_fn: Callable,
device: str = "cuda",
)
| 328 | |
| 329 | |
| 330 | def run_stability( |
| 331 | model_new, |
| 332 | get_inputs_fn: Callable, |
| 333 | device: str = "cuda", |
| 334 | ) -> Dict[str, Any]: |
| 335 | """Test numerical stability: check for NaN/Inf on normal inputs.""" |
| 336 | import torch |
| 337 | |
| 338 | result: Dict[str, Any] = {"stability": "PASS", "details": []} |
| 339 | |
| 340 | for trial in range(3): |
| 341 | try: |
| 342 | inputs = get_inputs_fn() |
| 343 | inputs_dev = [ |
| 344 | inp.to(device) if isinstance(inp, torch.Tensor) else inp |
| 345 | for inp in inputs |
| 346 | ] |
| 347 | with torch.no_grad(): |
| 348 | output = model_new(*inputs_dev) |
| 349 | |
| 350 | if isinstance(output, torch.Tensor): |
| 351 | if _has_nan_inf(output): |
| 352 | result["stability"] = "WARN" |
| 353 | result["details"].append(f"trial {trial}: output contains NaN/Inf") |
| 354 | elif isinstance(output, (tuple, list)): |
| 355 | for i, o in enumerate(output): |
| 356 | if isinstance(o, torch.Tensor) and _has_nan_inf(o): |
| 357 | result["stability"] = "WARN" |
| 358 | result["details"].append(f"trial {trial}: output[{i}] contains NaN/Inf") |
| 359 | |
| 360 | except Exception as e: |
| 361 | result["stability"] = "FAIL" |
| 362 | result["details"].append(f"trial {trial}: {type(e).__name__}: {e}") |
| 363 | |
| 364 | return result |
| 365 | |
| 366 | |
| 367 | def run_determinism( |