From 3c296a610125fb1ffea205bda0c48b7dc332716b Mon Sep 17 00:00:00 2001 From: Aarav <211792885+AaravMehta-07@users.noreply.github.com> Date: Thu, 15 Oct 2026 21:14:57 +0530 Subject: [PATCH 1/2] Preserve float64 precision in bootstrap and sign-flip statistics Keep double-precision working tensors while retaining float32 promotion for half and bfloat16 inputs. Add exact-value regressions for bootstrap bounds and sign-flip p-values, plus dtype controls. --- tests/unit/utilities/test_statistics.py | 31 ++++++++++++++++++++++++ transformer_lens/utilities/statistics.py | 4 +-- 2 files changed, 33 insertions(+), 2 deletions(-) diff --git a/tests/unit/utilities/test_statistics.py b/tests/unit/utilities/test_statistics.py index a881890ab..c4cc11362 100644 --- a/tests/unit/utilities/test_statistics.py +++ b/tests/unit/utilities/test_statistics.py @@ -76,6 +76,37 @@ def test_bootstrap_ci_matches_per_cell_loop_and_median(): bootstrap_ci(x, confidence=1.0) +@pytest.mark.parametrize("statistic", ["mean", "median"]) +def test_bootstrap_ci_preserves_float64_values(statistic: str): + x = torch.full((2, 8), 1.00000001, dtype=torch.float64) + low, high = bootstrap_ci(x, statistic=statistic, n_resamples=32) + torch.testing.assert_close(low, x[:, 0], rtol=0, atol=0) + torch.testing.assert_close(high, x[:, 0], rtol=0, atol=0) + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32, torch.float64]) +def test_statistics_support_floating_dtypes(dtype: torch.dtype): + x = torch.tensor([1.0, 2.0], dtype=dtype) + low, high = bootstrap_ci(torch.ones(8, dtype=dtype), n_resamples=32) + assert low.dtype == high.dtype == dtype + assert low.item() == high.item() == 1.0 + assert sign_flip_permutation_pvalue(x).item() == 0.5 + + +@pytest.mark.parametrize("exact_max_n", [0, 12]) +def test_sign_flip_float64_is_invariant_to_positive_scale(exact_max_n: int): + effects = torch.tensor([1.0, 2.0], dtype=torch.float64) + expected = sign_flip_permutation_pvalue(effects, exact_max_n=exact_max_n) + actual = sign_flip_permutation_pvalue(effects * 1e-50, exact_max_n=exact_max_n) + assert expected.item() < 1.0 + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + +def test_sign_flip_preserves_small_float64_effects(): + effects = torch.tensor([1.0, 1e-8], dtype=torch.float64) + assert sign_flip_permutation_pvalue(effects).item() == 0.5 + + def test_sign_flip_exact_enumeration_matches_hand_count(): effects = torch.tensor([[1.0, 2.0, 3.0, -0.5]]) observed = effects.mean().abs() diff --git a/transformer_lens/utilities/statistics.py b/transformer_lens/utilities/statistics.py index 1e9b19072..754634098 100644 --- a/transformer_lens/utilities/statistics.py +++ b/transformer_lens/utilities/statistics.py @@ -88,7 +88,7 @@ def bootstrap_ci( idx = torch.randint(0, n, (count, n), generator=gen).to(moved.device) resampled = moved[..., idx] # [..., count, n] stats.append(reduce(resampled, dim=-1)) - draws = torch.cat(stats, dim=-1).float() # [..., n_resamples] + draws = torch.cat(stats, dim=-1).to(torch.promote_types(x.dtype, torch.float32)) alpha = (1.0 - confidence) / 2.0 low = torch.quantile(draws, alpha, dim=-1) high = torch.quantile(draws, 1.0 - alpha, dim=-1) @@ -118,7 +118,7 @@ def sign_flip_permutation_pvalue( raise ValueError(f"alternative must be one of {_ALTERNATIVES}, got {alternative!r}") if n_permutations < 1: raise ValueError("n_permutations must be >= 1") - moved = _to_last(effects, dim).float() + moved = _to_last(effects, dim).to(torch.promote_types(effects.dtype, torch.float32)) n = moved.shape[-1] if n < 1: raise ValueError("cannot permute an empty axis") From 6a833221e6a701b1d1b29ce02ca0d3f488ef72c0 Mon Sep 17 00:00:00 2001 From: Aarav <211792885+AaravMehta-07@users.noreply.github.com> Date: Fri, 16 Oct 2026 19:51:14 +0530 Subject: [PATCH 2/2] Count floating-point ties in sign-flip p-values Use a directional relative roundoff threshold without an absolute floor. Cover rounded ties, tiny-scale separation, and p-value dtype retention. Assisted-by: OpenAI Codex --- tests/unit/utilities/test_statistics.py | 48 +++++++++++++++++++++++- transformer_lens/utilities/statistics.py | 8 ++-- 2 files changed, 51 insertions(+), 5 deletions(-) diff --git a/tests/unit/utilities/test_statistics.py b/tests/unit/utilities/test_statistics.py index c4cc11362..dffac5cab 100644 --- a/tests/unit/utilities/test_statistics.py +++ b/tests/unit/utilities/test_statistics.py @@ -90,7 +90,9 @@ def test_statistics_support_floating_dtypes(dtype: torch.dtype): low, high = bootstrap_ci(torch.ones(8, dtype=dtype), n_resamples=32) assert low.dtype == high.dtype == dtype assert low.item() == high.item() == 1.0 - assert sign_flip_permutation_pvalue(x).item() == 0.5 + pvalue = sign_flip_permutation_pvalue(x) + assert pvalue.dtype == torch.promote_types(dtype, torch.float32) + assert pvalue.item() == 0.5 @pytest.mark.parametrize("exact_max_n", [0, 12]) @@ -104,7 +106,49 @@ def test_sign_flip_float64_is_invariant_to_positive_scale(exact_max_n: int): def test_sign_flip_preserves_small_float64_effects(): effects = torch.tensor([1.0, 1e-8], dtype=torch.float64) - assert sign_flip_permutation_pvalue(effects).item() == 0.5 + p = sign_flip_permutation_pvalue(effects) + assert p.dtype == torch.float64 + assert p.item() == 0.5 + + +def test_sign_flip_counts_float64_ties_after_rounding(): + effects = torch.tensor([-1.0, 0.4, -0.5, 0.6], dtype=torch.float64) + p = sign_flip_permutation_pvalue(effects) + assert p.dtype == torch.float64 + assert p.item() == 0.875 + + scaled_p = sign_flip_permutation_pvalue(effects * 1e-50) + assert scaled_p.dtype == torch.float64 + assert scaled_p.item() == 0.875 + + +@pytest.mark.parametrize("alternative, expected", [("greater", 11 / 16), ("less", 7 / 16)]) +def test_sign_flip_counts_float64_ties_for_one_sided_alternatives( + alternative: str, expected: float +): + effects = torch.tensor([-1.0, 0.4, -0.5, 0.6], dtype=torch.float64) + p = sign_flip_permutation_pvalue(effects, alternative=alternative) + assert p.dtype == torch.float64 + assert p.item() == expected + + +def test_sign_flip_does_not_count_distinct_means_as_ties_near_zero(): + effects = torch.tensor([1.0, -1e-8], dtype=torch.float64) + for scale in (1.0, 1e-50): + scaled_effects = effects * scale + greater = sign_flip_permutation_pvalue(scaled_effects, alternative="greater") + less = sign_flip_permutation_pvalue(scaled_effects, alternative="less") + assert greater.dtype == less.dtype == torch.float64 + assert greater.item() == 0.5 + assert less.item() == 0.75 + + +def test_sign_flip_monte_carlo_preserves_float64_pvalue_dtype(): + effects = torch.ones(13, dtype=torch.float64) + p = sign_flip_permutation_pvalue( + effects, n_permutations=64, generator=derive_generator(23), exact_max_n=12 + ) + assert p.dtype == torch.float64 def test_sign_flip_exact_enumeration_matches_hand_count(): diff --git a/transformer_lens/utilities/statistics.py b/transformer_lens/utilities/statistics.py index 754634098..92a8f9471 100644 --- a/transformer_lens/utilities/statistics.py +++ b/transformer_lens/utilities/statistics.py @@ -144,10 +144,12 @@ def sign_flip_permutation_pvalue( def _count_hits(null: torch.Tensor, observed: torch.Tensor, alternative: str) -> torch.Tensor: + # Relative to the observed statistic: tolerate reduction-order ties without an absolute floor. + tolerance = observed.abs() * (100 * torch.finfo(null.dtype).eps) if alternative == "two-sided": - hit = null.abs() >= observed.abs() + hit = null.abs() >= observed.abs() - tolerance elif alternative == "greater": - hit = null >= observed + hit = null >= observed - tolerance else: - hit = null <= observed + hit = null <= observed + tolerance return hit.sum(dim=-1).to(null.dtype)