Add jax_array coverage to debug_nans_test

PiperOrigin-RevId: 478079509
google · Sep 30, 2022 · fb8558c · fb8558c
1 parent ec41de2
commit fb8558c
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 6 deletions.
diff --git a/tests/BUILD b/tests/BUILD
@@ -81,13 +81,9 @@ jax_test(
     srcs = ["custom_object_test.py"],
 )
 
-py_test(
+jax_test(
     name = "debug_nans_test",
     srcs = ["debug_nans_test.py"],
-    deps = [
-        "//jax",
-        "//jax:test_util",
-    ],
 )
 
 py_test(

diff --git a/tests/debug_nans_test.py b/tests/debug_nans_test.py
@@ -150,7 +150,7 @@ def testPjit(self):
     if jax.device_count() < 2:
       raise SkipTest("test requires >=2 devices")
 
-    p = jax.experimental.PartitionSpec('x')
+    p = pjit.PartitionSpec('x')
     f = pjit.pjit(lambda x: 0. / x,
                   in_axis_resources=p,
                   out_axis_resources=p)