Configure tests to explicitly use jax_threefry_partitionable=False.

See jax-ml/jax#18480

PiperOrigin-RevId: 746173050

Author: jburnim

trax/layers/research/efficient_attention_test.py

@@ -25,6 +25,8 @@
 from trax.fastmath import numpy as jnp
 from trax.layers.research import efficient_attention
 
+jax.config.update('jax_threefry_partitionable', False)
+
 
 class EfficientAttentionTest(test.TestCase, parameterized.TestCase):