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

Reenable a GMRES test on CPU.

Browse Source 
It appears to pass once again with jaxlib 0.1.67. Unfortunately it just missed the 0.1.66 release.
This commit is contained in:
Peter Hawkins 2021-05-11 21:43:15 -04:00
parent 0ca259f7c2
commit 73df92bbad
1 changed files with 4 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
tests/lax_scipy_sparse_test.py
View File

@ -317,13 +317,15 @@ class LaxBackedScipyTests(jtu.JaxTestCase):
for dtype in [np.float64, np.complex128]
for preconditioner in [None, 'identity', 'exact', 'random']
for solve_method in ['incremental', 'batched']))
# TODO(b/186133663): test fails on CPU after LLVM change.
@jtu.skip_on_devices("cpu")
def test_gmres_against_scipy(
self, shape, dtype, preconditioner, solve_method):
if not config.x64_enabled:
raise unittest.SkipTest("requires x64 mode")
# The LLVM bug that caused this appears to be fixed in jaxlib 0.1.67.
if jtu.device_under_test() == "cpu" and jax.lib.version <= (0, 1, 66):
raise unittest.SkipTest("test fails on CPU jaxlib <= 0.1.66")
rng = jtu.rand_default(self.rng())
A = rng(shape, dtype)
b = rng(shape[:1], dtype)