Repository navigation
fix: keep Box-Muller randn finite by drawing u1 from (0, 1] - #1455
Merged
Merged
Conversation
All four Box-Muller blocks drew u1 = random.nextFloat() in [0, 1); a zero draw (P ~ 2^-24 per pair) makes ln(u1) = -Inf and injects Inf/NaN into the tensor. Across the hundreds of millions of samples a large model init makes this is near-certain - observed as all-NaN logits when initializing a Gemma 3 270M model from scratch. Draw u1 = 1 - nextDouble(), which lies in (0, 1], keep the transform in double precision until the final toFloat(), and hoist the shared sqrt term. Regression tests force the zero draw with a zero-bits Random and add a seeded mean/std sanity check. Fixes #1454 Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
|
📖 Documentation Preview The documentation has been built successfully for this PR. Generated Files:
Artifacts:
This comment will be updated automatically when the PR is updated. |
aharakal
approved these changes
Oct 31, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Fixes #1454.
What
All four Box-Muller
randnblocks drewu1 = random.nextFloat() ∈ [0, 1). A zero draw (P ≈ 2⁻²⁴ per pair) makesln(u1) = -Infand injectsInf/NaNinto the tensor — near-certain across the hundreds of millions of samples a large model init makes. Found in practice as all-NaN logits when initializing a Gemma 3 270M model from scratch on 0.62.4.How
u1 = 1.0 - random.nextDouble(), which lies in (0, 1] —ln(1) = 0is harmless.toFloat()(the old code roundedsqrt/costo float separately before multiplying).sqrt(-2 ln u1)term is hoisted in the pair variants (it was computed twice).Touched sites:
DenseTensorDataFactory— single-sample overload + theTensorDataFactory.randnoverride (FP32 and FP16/BF16 branches)MemorySegmentTensorDataFactory.randn(jvmMain)No public API change (no
apiDumpneeded).Tests
DenseTensorDataFactoryRandnTest(commonTest, kotlin.test): a zero-bitsRandomforces the worst-case draw for FP32/FP16/BF16 and the single-sample overload; plus a seeded 100k-sample mean/std check.MemorySegmentTensorDataFactoryRandnTest(jvmTest): same worst-case draw on the MemorySegment factory.Red/green verified: with the production fix stashed, 5 of the 6 new tests fail (the seeded statistical test passes either way, as expected); with the fix, all pass.
🤖 Generated with Claude Code