Increase matmul precision in tests to make them less flaky.

Author: nshepperd

tests/test_cross.py

@@ -11,6 +11,7 @@
 import numpy as np
 import math
 import einops
+jax.config.update("jax_default_matmul_precision", "highest")
 
 from flash_attn_jax import flash_mha
 from .ref_mha import ref_mha

tests/test_flash.py

@@ -12,6 +12,7 @@
 import numpy as np
 import math
 import einops
+jax.config.update("jax_default_matmul_precision", "highest")
 
 from flash_attn_jax import flash_mha
 from .ref_mha import ref_mha

tests/test_ring.py

@@ -17,6 +17,7 @@
 from functools import partial
 import einops
 import math
+jax.config.update("jax_default_matmul_precision", "highest")
 
 from flash_attn_jax.ring_attention import ring_fwd, ring_bwd
 from .ref_mha import ref_fwd, ref_bwd

tests/test_sharding.py

@@ -14,6 +14,7 @@
 from jax.sharding import PartitionSpec as P
 from jax.tree_util import tree_map
 jax.config.update("jax_traceback_filtering", "off")
+jax.config.update("jax_default_matmul_precision", "highest")
 
 from flash_attn_jax import flash_mha
 from .ref_mha import ref_mha

tests/test_varlen.py

@@ -12,6 +12,7 @@
 import numpy as np
 import math
 import einops
+jax.config.update("jax_default_matmul_precision", "highest")
 
 from flash_attn_jax import flash_mha
 from flash_attn_jax.varlen import flash_mha_varlen