Allow unit tests to disable 64 bit precision mode in jax.

PiperOrigin-RevId: 550617289

Author: ursk

tensorflow_probability/python/internal/test_util.py

@@ -2010,9 +2010,9 @@ def getTestCaseNames(self, testCaseClass):  # pylint:disable=invalid-name
     return names
 
 
-def main(jax_mode=JAX_MODE):
+def main(jax_mode=JAX_MODE, jax_enable_x64=True):
   """Test main function that injects a custom loader."""
-  if jax_mode:
+  if jax_mode and jax_enable_x64:
     from jax.config import config  # pylint: disable=g-import-not-at-top
     config.update('jax_enable_x64', True)