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

Relax some test tolerances to fix failures on Linux aarch64.

Browse Source 
PiperOrigin-RevId: 565930178
This commit is contained in:
Peter Hawkins 2023-09-16 06:54:52 -07:00 committed by jax authors
parent bf40f75bd5
commit f863cfbaad
3 changed files with 3 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
tests/BUILD
View File

@ -662,6 +662,7 @@ jax_test(
name = "pytorch_interoperability_test",
srcs = ["pytorch_interoperability_test.py"],
disable_backends = ["tpu"],
tags = ["not_build:arm"],
deps = py_deps("torch"),
)

2
tests/lax_numpy_test.py
View File

@ -1278,7 +1278,7 @@ class LaxBackedNumpyTests(jtu.JaxTestCase):
elif rank == 2 and jtu.device_under_test() in ("tpu", "gpu"):
self.skipTest("Nonsymmetric eigendecomposition is only implemented on the CPU backend.")
rng = jtu.rand_default(self.rng())
tol = { np.int8: 1e-3, np.int32: 1e-3, np.float32: 1e-3, np.float64: 1e-6 }
tol = { np.int8: 2e-3, np.int32: 1e-3, np.float32: 1e-3, np.float64: 1e-6 }
if jtu.device_under_test() == "tpu":
tol[np.int32] = tol[np.float32] = 1e-1
tol = jtu.tolerance(dtype, tol)

2
tests/lax_scipy_test.py
View File

@ -489,7 +489,7 @@ class LaxBackedScipyTests(jtu.JaxTestCase):
should_be_eye = np.matmul(unitary.conj().T, unitary)
else:
should_be_eye = np.matmul(unitary, unitary.conj().T)
tol = 500 * float(jnp.finfo(matrix.dtype).eps)
tol = 650 * float(jnp.finfo(matrix.dtype).eps)
eye_mat = np.eye(should_be_eye.shape[0], dtype=should_be_eye.dtype)
with self.subTest('Test unitarity.'):
self.assertAllClose(