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

Update test since C PjRt API topology description is available.

Browse Source 
The compilation_cache_test had an exclusion since the C PjRt
topology description had not been implemented. Now that it is
available, remove the exclusion.

PiperOrigin-RevId: 581396824
This commit is contained in:
jax authors 2023-11-10 16:11:35 -08:00
parent 11236dbe34
commit 45982c8439
1 changed files with 0 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
tests/compilation_cache_test.py
View File

@ -331,9 +331,6 @@ class CompilationCacheTest(jtu.JaxTestCase):
"/jax/compilation_cache/original_compile_time_saved_sec",
durations)
else:
# TODO(b/293308239) Remove this skipping when pjrt c api is supported.
if xla_bridge.using_pjrt_c_api():
raise SkipTest("PJRT C API not supported yet.")
self.assertNotIn(
"/jax/compilation_cache/compile_time_saved_sec", durations)
@ -427,9 +424,6 @@ class CompilationCacheTest(jtu.JaxTestCase):
- previous_counts["/jax/compilation_cache/cache_hits_original"],
1)
else:
# TODO(b/293308239) Remove this skipping when pjrt c api is supported.
if xla_bridge.using_pjrt_c_api():
raise SkipTest("PJRT C API not supported yet.")
self.assertEqual(
_counts["/jax/compilation_cache/cache_hits"]
- previous_counts["/jax/compilation_cache/cache_hits"],