Configure tests to explicitly use jax_threefry_partitionable=False.

See jax-ml/jax#18480

PiperOrigin-RevId: 746378505

Author: jburnim

discussion/pathfinder/pathfinder_test.py

@@ -19,6 +19,8 @@
 from discussion.pathfinder import pathfinder
 import tensorflow_probability.substrates.jax as tfp
 
+jax.config.update("jax_threefry_partitionable", False)
+
 tfd = tfp.distributions
 
 

spinoffs/autobnn/autobnn/bnn_tree_test.py

@@ -22,6 +22,8 @@
 from autobnn import kernels
 from absl.testing import absltest
 
+jax.config.update('jax_threefry_partitionable', False)
+
 
 class TreeTest(parameterized.TestCase):
 

spinoffs/autobnn/autobnn/kernels_test.py

@@ -24,6 +24,8 @@
 
 from absl.testing import absltest
 
+jax.config.update('jax_threefry_partitionable', False)
+
 
 KERNELS = [
     kernels.IdentityBNN,

tensorflow_probability/python/experimental/fastgp/fast_gp_test.py

@@ -821,4 +821,5 @@ class FastGpTestFloat64(_FastGpTest):
 
 if __name__ == "__main__":
   config.update("jax_enable_x64", True)
+  config.update("jax_threefry_partitionable", False)
   absltest.main()

tensorflow_probability/python/experimental/fastgp/fast_log_det_test.py

@@ -25,6 +25,8 @@
 from tensorflow_probability.substrates import jax as tfp
 from tensorflow_probability.substrates.jax.internal import test_util
 
+jax.config.update('jax_threefry_partitionable', False)
+
 # pylint: disable=invalid-name
 
 

tensorflow_probability/python/experimental/fastgp/fast_mtgp_test.py

@@ -458,4 +458,5 @@ class FastMultiTaskGpTestFloat64(_FastMultiTaskGpTest):
 
 if __name__ == "__main__":
   config.update("jax_enable_x64", True)
+  config.update("jax_threefry_partitionable", False)
   absltest.main()

tensorflow_probability/python/experimental/fastgp/preconditioners_test.py

@@ -696,4 +696,5 @@ class PreconditionersTestFloat64(_PreconditionersTest):
 
 if __name__ == "__main__":
   config.update("jax_enable_x64", True)
+  config.update("jax_threefry_partitionable", False)
   absltest.main()