Combine bits via xor when bitcasting from larger to smaller type#734
Merged
patrick-kidger merged 3 commits intopatrick-kidger:mainfrom Feb 23, 2026
Conversation
Previously, when casting e.g. a float64 to an int32, numbers that were close in float64 could be mapped to identical int32's. Since these int's are used as keys to generate random sequences, this is problematic, as it results in identical noise being generated in subsequent timesteps. This commit fixes this by not throwing away bits when the input type is larger than the requested output type. Instead, the larger number is bitcast to multiple values in the smaller type, which are then combined using xor.
Owner
|
Looks reasonable to me! (Though I think the pre-commits are failing, see |
Comment on lines
+21
to
+24
| if result.shape != val.shape: | ||
| result = jnp.bitwise_xor.reduce(result, axis=-1) | ||
|
|
||
| return result |
Owner
There was a problem hiding this comment.
Maybe an assert result.shape == val.shape afterwards to be sure that this xor has done its job?
This checks the intended behaviour of mapping nearby numbers to distinct values when downcasting to a smaller dtype.
Contributor
Author
|
Added two assertions (shape and dtype) and a simple test which fails with the previous implementation, but passes with these changes. The formatting issues have also been fixed. LMK if any further additions are necessary! |
Owner
|
Looks perfect! Merged 🎉 Thank you for the fix. |
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
2 participants
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 #733.
Previously, when downcasting to a type with fewer bits, the larger type would be converted to a smaller intermediate dtype, which was subsequently bitcasted to the requested type. The casting to a smaller intermediate type loses precision, and can result in distinct values being cast to identical values. Since these values are used as keys to generate random sequences, this is problematic, as it results in identical noise being generated in subsequent timesteps. This commit fixes this by not throwing away bits when the input type is larger than the requested output type. Instead, the larger number is bitcast to multiple values in the smaller type, which are then combined using xor.