Update tests/jax_utils_test.py

Commit

11 days ago

Update tests/jax_utils_test.py Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>