Remove tests for zeros_like and add_any primitives from neural_tangents.

These primitives have been removed from JAX at head, and these tests will fail with a current JAX. PiperOrigin-RevId: 595974695
google · Feb 1, 2024 · f96f176 · f96f176
1 parent ad47437
commit f96f176
Showing 1 changed file with 0 additions and 6 deletions.
diff --git a/tests/rules_test.py b/tests/rules_test.py
@@ -262,8 +262,6 @@ def _concat_shapes(max_n_args: int = 4, *shapes):
 
  jax.lax.copy_p: lambda s, _: [{}],
 
- ad.zeros_like_p: lambda s, _: [{}],
-
  lax.neg_p: lambda s, _: [{}],
 
  lax.transpose_p: lambda s, _: [
@@ -372,10 +370,6 @@ def _concat_shapes(max_n_args: int = 4, *shapes):
 
 
 _BINARY_PRIMITIVES = {
- # TODO(romann): what is the purpose of this primitive?
- ad.add_jaxvals_p:
- lambda s1, s2: ([{}] if s1 == s2 else []),
-
  lax.mul_p:
  lambda s1, s2: ([{}] if _is_broadcastable(s1, s2) else []),