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

Disable lax_scipy_sparse_test.py cases that are hanging on GPU.

Browse Source 
See #6471.
This commit is contained in:
Skye Wanderman-Milne 2021-04-15 15:54:15 -07:00
parent f65a327c76
commit 346df9c557
1 changed files with 4 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
tests/lax_scipy_sparse_test.py
View File

@ -253,6 +253,7 @@ class LaxBackedScipyTests(jtu.JaxTestCase):
for dtype in float_types + complex_types
for preconditioner in [None, 'identity', 'exact']
))
@jtu.skip_on_devices("gpu")
def test_bicgstab_on_identity_system(self, shape, dtype, preconditioner):
A = jnp.eye(shape[1], dtype=dtype)
solution = jnp.ones(shape[1], dtype=dtype)
@ -277,6 +278,7 @@ class LaxBackedScipyTests(jtu.JaxTestCase):
for dtype in float_types + complex_types
for preconditioner in [None, 'identity', 'exact']
))
@jtu.skip_on_devices("gpu")
def test_bicgstab_on_random_system(self, shape, dtype, preconditioner):
rng = jtu.rand_default(self.rng())
A = rng(shape, dtype)
@ -365,6 +367,7 @@ class LaxBackedScipyTests(jtu.JaxTestCase):
for preconditioner in [None, 'identity', 'exact']
for solve_method in ['batched', 'incremental']
))
@jtu.skip_on_devices("gpu")
def test_gmres_on_identity_system(self, shape, dtype, preconditioner,
solve_method):
A = jnp.eye(shape[1], dtype=dtype)
@ -395,6 +398,7 @@ class LaxBackedScipyTests(jtu.JaxTestCase):
for preconditioner in [None, 'identity', 'exact']
for solve_method in ['incremental', 'batched']
))
@jtu.skip_on_devices("gpu")
def test_gmres_on_random_system(self, shape, dtype, preconditioner,
solve_method):
rng = jtu.rand_default(self.rng())