Avoid overflow in attribution variance
This commit is contained in:
@@ -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(
|
||||
|
||||
Reference in New Issue
Block a user