diff --git a/tensorflow_probability/python/internal/test_util.py b/tensorflow_probability/python/internal/test_util.py
index 67af7e0ec1..9af934ed5f 100644
--- a/tensorflow_probability/python/internal/test_util.py
+++ b/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)