Configure tests to explicitly use jax_threefry_partitionable=False.

See jax-ml/jax#18480

PiperOrigin-RevId: 746157829

Author: jburnim

bayeux/tests/mcmc_test.py

@@ -22,6 +22,8 @@
 import pytest
 import tensorflow_probability.substrates.jax as tfp
 
+jax.config.update("jax_threefry_partitionable", False)
+
 tfd = tfp.distributions
 
 METHODS = [getattr(bx.mcmc, k).name for k in bx.mcmc.__all__]