fix a few test bitwise op tests for np 2.0

PiperOrigin-RevId: 672999286
google · Sep 10, 2024 · 17c19c7 · 17c19c7
1 parent 39c9bbb
commit 17c19c7
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/trax/tf_numpy/jax_tests/lax_numpy_test.py b/trax/tf_numpy/jax_tests/lax_numpy_test.py
@@ -272,13 +272,13 @@ def minus(a, b):
 ]
 
 JAX_BITWISE_OP_RECORDS = [
-    op_record("bitwise_and", 2, int_dtypes + unsigned_dtypes, all_shapes,
+    op_record("bitwise_and", 2, int_dtypes, all_shapes,
               jtu.rand_default, []),
-    op_record("bitwise_not", 1, int_dtypes + unsigned_dtypes, all_shapes,
+    op_record("bitwise_not", 1, int_dtypes, all_shapes,
               jtu.rand_default, []),
-    op_record("bitwise_or", 2, int_dtypes + unsigned_dtypes, all_shapes,
+    op_record("bitwise_or", 2, int_dtypes, all_shapes,
               jtu.rand_default, []),
-    op_record("bitwise_xor", 2, int_dtypes + unsigned_dtypes, all_shapes,
+    op_record("bitwise_xor", 2, int_dtypes, all_shapes,
               jtu.rand_default, []),
 ]