test: compare same-device regression round-trips at float32 precision

[ggprior]

main fails two slow MPS save/load tests on the macos-15 job:

• test__save_and_load_fit_with_cache__predictions_equal[regression-mps-mps-v3]
• test__save_and_load_fit_with_cache_twice__predictions_equal[regression-mps-mps-v3]

One of 40 regression predictions differs by 1.5e-5 (8.6e-8 relative) between the original and the reloaded model, which is one float32 step at a value of about 179. The check used assert_array_almost_equal, which requires agreement to 6 decimals. For values in the hundreds that means bit-identical float32 outputs, which float32 inference doesn't guarantee.

Change

Same-device regression round-trips now compare with assert_allclose(atol=1e-6 × max(|prediction|, 1)). That's float32-level precision relative to the predictions' own scale; for these predictions it allows about 2.5e-4. Classifier labels and probabilities keep the existing check.

Since #1311, regression outputs are decoded in z-space and mapped back in float64, which moved where float32 rounding happens and exposed the existing difference on one row. The underlying difference comes from the regression border remap. On non-CUDA devices it accumulates with index_add_, which on MPS sums in a different order on each call, so the aggregated logits can change in the last bit between two predicts of the same model. Everything upstream (weights, KV caches, per-estimator outputs) is bit-identical. A deterministic remap on MPS is a separate change.

Verification

PR CI skips slow tests. Two temporary commits made this PR's macos-15 job run tests/test_save_load_fitted_model.py including the slow MPS cases, then repeat the four MPS fit_with_cache regression round-trips 10 times. All 41 + 10 × 4 passed on the fix (job, along with the rest of CI. After rebasing onto current main, a further temporary commit ran the whole suite including slow tests on macos-15, the same scope as the post-merge main run: 2646 passed (job). The temporary commits have been dropped; the failure is intermittent (a CI-only run without the fix also happened to pass), which is the nondeterminism described above.

🤖 Generated with Claude Code

https://claude.ai/code/session_018qH4DKrmQVSS3qbzViv5nf

[oscarkey]

Just pushed a change to compute the atol excluding the shift. z*scale + shift happens in float64, so we I think we should only consider the error in z, scaled up by scale. If we include the shift, the tolerance can end up being too large? In practice, the test data is centred, so shift is close to zero, but just in case.

[jmkuebler]

lgtm