Add comprehensive test suite for random_u48 function#169
Open
oliverdutton wants to merge 2 commits into
Open
Conversation
This commit adds a new test class TestRandomU48 with 5 test cases: 1. test_random_u48_basic: Verifies basic functionality and range checking 2. test_random_u48_with_external_callback: Uses jax.pure_callback with arbitrary precision Python integers to verify the correctness of the modulo operation 3. test_random_u48_distribution: Tests that the function produces varied values across multiple seeds 4. test_random_u48_batched: Tests batch processing with multiple indices 5. test_random_u48_large_max_val: Tests with large values close to U48 maximum (2^47) The tests use jax.random.key() and jax.random.split() to create the 4 keys needed for u128 random generation, as suggested in the task description. https://claude.ai/code/session_013SJnH3GeqezZzM7V8tJkTg
Updated all test cases to use 2D array shapes for indices and max_val, matching the actual usage in topp_mask.py and topk_topp_mask_and_sample.py. Changes: - Changed all dim0_indices from 1D to 2D shape (e.g., [[0], [1], [2]]) - Changed max_val initialization to use 2D arrays (e.g., [[1000000]]) - Updated array indexing to account for 2D shapes - Renamed test_random_u48_with_external_callback to test_random_u48_with_high_precision - Simplified high precision test to use Python's arbitrary precision integers directly instead of jax.pure_callback (which doesn't support object dtype) - Adjusted max values to fit within int32 range (2^30 and 2^31-1 instead of 2^40 and 2^47) All 5 tests now pass successfully: 1. test_random_u48_basic - Basic functionality with multiple indices 2. test_random_u48_with_high_precision - Python arbitrary precision verification 3. test_random_u48_distribution - Distribution quality across 100 samples 4. test_random_u48_batched - Batch processing with 10 values 5. test_random_u48_large_max_val - Large max values (2^31-1) https://claude.ai/code/session_013SJnH3GeqezZzM7V8tJkTg
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.
Summary
This PR adds a complete test suite for the
random_u48function with five test cases covering basic functionality, high-precision correctness, distribution properties, batched operations, and large value handling.Key Changes
TestRandomU48test class with five comprehensive test methods:test_random_u48_basic: Validates basic random U48 generation and range checkingtest_random_u48_with_high_precision: Verifies correctness using Python's arbitrary precision integers for values up to 2^30test_random_u48_distribution: Ensures generated values have sufficient variety across 100 samplestest_random_u48_batched: Tests batched operation with multiple indices and validates shape/rangetest_random_u48_large_max_val: Tests with maximum int32 values (2^31 - 1) to verify large value handlingrandom_u48function into test moduleTestRandomU48in the test runnerNotable Implementation Details
[batch_size, 1]) to match the expected usage pattern intopp_mask.py[0, max_val)https://claude.ai/code/session_013SJnH3GeqezZzM7V8tJkTg