Skip to content

Extend test coverage for jax_utils.py#2023

Closed
irhyl wants to merge 1 commit intogoogle-deepmind:mainfrom
gsoc-2026:pr/jax-utils-test-coverage
Closed

Extend test coverage for jax_utils.py#2023
irhyl wants to merge 1 commit intogoogle-deepmind:mainfrom
gsoc-2026:pr/jax-utils-test-coverage

Conversation

@irhyl
Copy link
Copy Markdown

@irhyl irhyl commented Mar 1, 2026

Summary

Doubles the test coverage for jax_utils.py from 15 to 32 tests by adding coverage for get_np_dtype, env_bool, assert_rank, and error_if.

Changes

  • Modified: torax/_src/tests/jax_utils_test.py
    • 4 tests for get_np_dtype (default, f64, f32, invalid dtype)
    • 8 tests for env_bool (true/false/1/0 strings, case insensitivity, missing env var, invalid)
    • 3 tests for assert_rank (correct, wrong, scalar)
    • 2 tests for error_if (passthrough behavior)
    • Added get_np_dtype.cache_clear() in setUp for test isolation

Testing

All 32 tests pass (15 pre-existing + 17 new).

Add 17 new tests covering previously untested functions:
- get_np_dtype: default/f64/f32 dtype selection, invalid dtype error
- env_bool: true/false/1/0 strings, case insensitivity, missing env
  var default, invalid value raises
- assert_rank: correct rank, wrong rank error, scalar rank
- error_if: passthrough behavior when errors disabled

Also add cache_clear() in setUp to ensure test isolation for
cached dtype functions.
@Nush395
Copy link
Copy Markdown
Collaborator

Nush395 commented Mar 27, 2026

Thanks for the contribution. I think only the tests on the np dtype are worth adding here. The rest of the tests are checking logic on error_if and assert rank that I don't think adds much value unfortunately. Please feel free to reopen a smaller PR with only the np tests but otherwise I'm going to close this for now.

@Nush395 Nush395 closed this Mar 27, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants