Avoid overflow in attribution variance

This commit is contained in:
2026-07-01 10:58:09 +08:00
parent 16df74c04d
commit d3e68ac09b

View File

@@ -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(