From d3e68ac09bfa157b25f9aa31272560fa391654c0 Mon Sep 17 00:00:00 2001 From: Jiarui Li Date: Wed, 1 Jul 2026 10:58:09 +0800 Subject: [PATCH] Avoid overflow in attribution variance --- evaluate_extra_info_attribution.py | 16 ++++++++++------ 1 file changed, 10 insertions(+), 6 deletions(-) diff --git a/evaluate_extra_info_attribution.py b/evaluate_extra_info_attribution.py index cb8c425..818c31b 100644 --- a/evaluate_extra_info_attribution.py +++ b/evaluate_extra_info_attribution.py @@ -364,6 +364,11 @@ def iter_extra_info_ablated_batches( yield concat_tensor_batches(pending_batches), pending_types, pending_rows +def finite_float64(values: Any) -> np.ndarray: + arr = np.asarray(values, dtype=np.float64) + return arr[np.isfinite(arr)] + + def update_death_summary( summary: dict[tuple[Any, ...], dict[str, float]], *, @@ -390,10 +395,10 @@ def update_death_summary( ) acc["n"] += float(len(group)) for column in DEATH_PARAMETER_COLUMNS: - vals = pd.to_numeric(group[column], errors="coerce").dropna() - acc[f"count__{column}"] += float(len(vals)) + vals = finite_float64(pd.to_numeric(group[column], errors="coerce")) + acc[f"count__{column}"] += float(vals.size) acc[f"sum__{column}"] += float(vals.sum()) - acc[f"sumsq__{column}"] += float((vals * vals).sum()) + acc[f"sumsq__{column}"] += float(np.square(vals).sum()) def update_disease_parameter_summary( @@ -425,11 +430,10 @@ def update_disease_parameter_summary( ) acc["n"] += float(vals_2d.shape[0]) for col_idx, column in enumerate(DISEASE_PARAMETER_COLUMNS): - vals = vals_2d[:, col_idx] - vals = vals[np.isfinite(vals)] + vals = finite_float64(vals_2d[:, col_idx]) acc[f"count__{column}"] += float(vals.size) acc[f"sum__{column}"] += float(vals.sum()) - acc[f"sumsq__{column}"] += float((vals * vals).sum()) + acc[f"sumsq__{column}"] += float(np.square(vals).sum()) def write_death_summary_csv(