Attempt to fix float64 test issue on GitHub Actions.
Attempt to fix float64 test issue on GitHub Actions.

I can't reproduce this on my corp workstation, personal machine (Ubuntu in WSL)
but a hypothesis is that for some reason the version of pytest in the GitHub
action is not running the module as main, so the side effect of updating
jax.config is not happening. Doing this as part of setUp/tearDown should ensure
it happens as part of running the test suite.

PiperOrigin-RevId: 351350381
Change-Id: Id3f79f7fcc4991986714d1f7e2ab146083fccc9c

deepmind/dm-haiku,deepmind/dm-haiku