Configure tests to explicitly use jax_threefry_partitionable=False.

See jax-ml/jax#18480

PiperOrigin-RevId: 746182723

Author: jburnim

mctx/_src/tests/policies_test.py

@@ -22,6 +22,8 @@
 from mctx._src import policies
 import numpy as np
 
+jax.config.update("jax_threefry_partitionable", False)
+
 
 def _make_bandit_recurrent_fn(rewards, dummy_embedding=()):
   """Returns a recurrent_fn with discount=0."""

mctx/_src/tests/tree_test.py

@@ -26,6 +26,8 @@
 import mctx
 import numpy as np
 
+jax.config.update("jax_threefry_partitionable", False)
+
 
 def _prepare_root(batch_size, num_actions):
   """Returns a root consistent with the stored expected trees."""