Deflake tests

PiperOrigin-RevId: 538511611
google · Aug 2, 2023 · 1e82970 · 1e82970
1 parent 0189579
commit 1e82970
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/tests/experimental/empirical_tf_test.py b/tests/experimental/empirical_tf_test.py
@@ -228,7 +228,7 @@ def _compare_ntks(
 
     x1_jax = np.array(x1)
     x2_jax = np.array(x2)
-    params_jax = jax.tree_map(lambda x: np.array(x), params)
+    params_jax = jax.tree_map(np.array, params)
 
     jax_ntks = [ntk_fn_i(x1_jax, x2_jax, params_jax)
                 for ntk_fn_i in jax_ntk_fns]

diff --git a/tests/rules_test.py b/tests/rules_test.py
@@ -61,9 +61,9 @@
     (),
     # (0,),
     (2,),
-    (0, 0),
-    (1, 0),
-    (0, 1),
+    # (0, 0),
+    # (1, 0),
+    # (0, 1),
     # (1, 1),
     # (2, 1),
     (1, 2),