jax

5c4dd597 - Use a np.printoptions context manager rather than np.set_printoptions in test utility.

Use a np.printoptions context manager rather than np.set_printoptions in test utility. I am guessing this may be responsible for the TPU CI flake seen here: https://github.com/jax-ml/jax/actions/runs/23086906115/job/67065075706 This would also explain why we have only ever seen this failure under pytest and never bazel: it requires that you run both a test using this utility and the DebugPrintTest test case in the same pytest-xdist worker. ``` _________ DebugPrintTest.test_debug_print_respects_numpy_printoptions __________ [gw0] linux -- Python 3.14.3 /usr/bin/python3.14 tests/debugging_primitives_test.py:270: in test_debug_print_respects_numpy_printoptions self.assertEqual(output(), "[1.23 2.35 0. ]\n") E AssertionError: E - [1.2345 2.3456 0.0000001] E + [1.23 2.35 0. ] ``` PiperOrigin-RevId: 883663987