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

attempt to fix CI failure (from #3845 test?) (#3846)

Browse Source
This commit is contained in:
Matthew Johnson 2020-07-23 20:59:12 -07:00 committed by GitHub
parent cc9528d97d
commit e2424e3b24
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 5 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
tests/api_test.py
View File

@ -2534,13 +2534,17 @@ class CustomJVPTest(jtu.JaxTestCase):
def run():
return jax.grad(unroll)(jnp.array(1.0), jnp.array([1.0]))
expected = run()
# we just don't want this to crash
n_workers = 20
with concurrent.futures.ThreadPoolExecutor(max_workers=n_workers) as e:
futures = []
for _ in range(n_workers):
futures.append(e.submit(run))
_ = [f.result() for f in futures]
results = [f.result() for f in futures]
for ans in results:
self.assertAllClose(ans, expected)
class CustomVJPTest(jtu.JaxTestCase):