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

Relax test tolerance to fix BCSR sparse matmul test failure on P100 GPU.

Browse Source 
PiperOrigin-RevId: 563441383
This commit is contained in:
Peter Hawkins 2023-09-07 08:36:51 -07:00 committed by jax authors
parent 429422dfea
commit 9b447aa3ec
1 changed files with 3 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
tests/sparse_test.py
View File

@ -1934,11 +1934,12 @@ class BCOOTest(sptu.SparseTestCase):
args_maker = lambda: [sprng(lhs_shape, lhs_dtype, n_batch=n_batch_lhs),
jnp.array(rng(rhs_shape, rhs_dtype))]
tol = {np.float64: 1E-7, np.complex128: 1E-7,
tol = {np.float64: 1E-7, np.complex128: 1E-6,
np.float32: 2E-6, np.complex64: 2E-6}
with jtu.strict_promotion_if_dtypes_match([lhs_dtype, rhs_dtype]):
self._CheckAgainstDense(operator.matmul, operator.matmul, args_maker, tol=tol)
self._CheckAgainstDense(operator.matmul, operator.matmul, args_maker,
tol=tol)
@jtu.sample_product(