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

[sparse] re-enable bcoo_spdot_general test

Browse Source
This commit is contained in:
Jake VanderPlas 2021-11-15 16:21:04 -08:00
parent be751d1dd6
commit e4291e0b49
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
tests/sparse_test.py
View File

@ -1045,7 +1045,7 @@ class BCOOTest(jtu.JaxTestCase):
self.assertAllClose(jr_sparse, jr_dense, atol=tol)
self.assertAllClose(jf_sparse, jr_sparse, atol=tol)
@unittest.skip("Jaxlib GPU build failing in OSS.")
@unittest.skipIf(jtu.device_under_test() == "tpu", "TPU has insufficient precision")
@parameterized.named_parameters(jtu.cases_from_list(
{"testcase_name": "_{}[n_batch={}]_{}[n_batch={}]_dims={}".format(
jtu.format_shape_dtype_string(lhs_shape, dtype), lhs_n_batch,