mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-17 20:36:05 +00:00
Code Issues Packages Projects Releases Wiki Activity

Fix an A100 nightly unit test failure on testDotGeneral() by replacing TF32 with float32

Browse Source
This commit is contained in:
Jane Liu 2023-09-12 14:30:40 -07:00
parent 801cbef011
commit 90da4f153a
1 changed files with 2 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
tests/batching_test.py
View File

@ -222,6 +222,8 @@ class BatchingTest(jtu.JaxTestCase):
self.assertAllClose(ans[i], expected_ans, check_dtypes=False)
# Replace the default TF32 with float32 in order to make it pass on A100
@jax.default_matmul_precision("float32")
def testDotGeneral(self):
R = self.rng().randn