Preserve float64 precision in bootstrap and sign-flip statistics

[AaravMehta-07]

Description

Preserve float64 precision in bootstrap_ci and sign_flip_permutation_pvalue, and count floating-point ties consistently in the permutation test.

• Use torch.promote_types(input.dtype, torch.float32) instead of narrowing intermediate tensors to float32. Half/bfloat16 still use float32 working tensors; float64 p-values remain float64, and bootstrap bounds retain the input dtype.
• Use directional comparisons with a relative roundoff threshold of 100 * eps * abs(observed), following SciPy's permutation-test convention. There is no absolute tolerance floor, so genuinely different tiny statistics are not merged into ties.
• Keep generator streams, resampling, chunking and public signatures unchanged. No dependencies added.

The narrowing can change conclusions: constant float64 samples of 1.00000001 lose precision; exact effects [1.0, 1e-8] incorrectly produce p=1 rather than 0.5; scaling by 1e-50 can underflow the effects. The review also exposed a separate tie issue: rounded float64 effects [-1.0, 0.4, -0.5, 0.6] gave p=0.75 instead of the exact-enumeration result 0.875. Both precision issues are now covered.

Type Of Change

• Bug fix

Validation

• Original narrowing regressions against unchanged dev: 5 failed, 11 passed before the dtype repair; the original related run passed 33 tests after it.
• New rounded-tie regression before this follow-up: failed, 0.75 versus expected 0.875; now passes, including scaling by 1e-50.
• Current CPU runs with normal repository pytest options, including jaxtyping/beartype instrumentation: 21 statistics tests, 106 sparse-probing tests, and the real patching-result statistics consumer passed.
• Added explicit p-value dtype assertions for supported floating dtypes and the float64 Monte Carlo path; one-sided ties and genuinely distinct tiny means are covered.
• Independent verification passed 45 cases, including integer-oracle comparisons, all alternatives, positive/negative and scaled effects, Monte Carlo draws, dtype promotion, zero and noncontiguous axes.
• make format: exit 0. Its existing isort format . recipe reports a missing format path, so both changed files were also checked separately with proper isort --check-only and black --check commands: passed. Notebook formatting was unavailable because Black's Jupyter extra is not installed; no notebooks changed.
• uv run ... python -m mypy .: no issues in 393 source files.
• git diff --check: passed.

Validation used actual source with CPU runtimes, including WSL Python 3.12. No full acceptance/integration, model-download, notebook, CUDA/MPS or performance suite was run. Initial local consumer runs accidentally disabled pytest addopts; their logs were retained, the runner was corrected, and the passing results above use the normal configuration. No tests or runtime contracts were weakened to address that runner error.

Checklist

• Failing-before regressions and floating-dtype controls.
• New and related existing tests pass locally with normal runtime instrumentation.
• No added skips, xfails, weakened assertions or type suppressions.
• Source changes limited to working-dtype promotion and permutation tie comparison.

Assistance And Scope

AI tools assisted with investigation, implementation, independent review and local testing. This is a CPU statistical-correctness fix; no GPU, pretrained-model parity, performance or model-quality result is claimed.

[jlarson4]

@AaravMehta-07 Thanks for putting this fix together, solid solution. One comment related to a unique case, and one comment pertaining to tightening our tests

[jlarson4]

Float64 narrows the rounding problem, but doesn't resolve it. _count_hits still compares with bare >=, and the observed and null means are summed in different orders. So sign patterns that tie exactly can still miss by an ulp and give an anti-conservative p-value, e.g. for effects rounded to one decimal. Could you make the comparison tolerant of floating-point ties, the way scipy's permutation_test does, and add a tie case to the tests?

[AaravMehta-07]

Thanks, both findings are addressed in 6a83322.

I reproduced the tie undercount with float64 effects [-1.0, 0.4, -0.5, 0.6]: the previous comparator returned 0.75 instead of the exact 0.875. _count_hits now uses SciPy's directional relative threshold, 100 * eps * abs(observed), without an absolute tolerance floor. Regressions cover two-sided/one-sided ties, scaling by 1e-50, and genuinely distinct tiny means. P-value dtype assertions now cover all supported floating inputs and the float64 Monte Carlo path.

With normal repository pytest options and runtime instrumentation: 21 statistics tests, 106 sparse-probing tests, and the real patching statistics consumer passed. Independent verification passed 45 numerical cases; full mypy passed all 393 source files. Formatting and whitespace checks passed; the existing Makefile isort-path warning and separate exact changed-file checks are documented in the updated description. The repair merges cleanly with current dev in a local merge-tree check. Hosted checks for this new head remain separate from these local results; no GPU/full acceptance qualification is claimed.

AI-assisted follow-up. I have left the review threads open for your assessment.

[jlarson4]

This checks the bootstrap dtype but not the p-value's, so a cast back to float32 on return would still pass. Can you please assert the p-value's dtype here too?