explicit tuple equality

xdslproject · Aug 17, 2024 · 34d1c59 · 34d1c59
1 parent 968bfa5
commit 34d1c59
Showing 1 changed file with 6 additions and 6 deletions.
diff --git a/tests/backend/test_jax_executable.py b/tests/backend/test_jax_executable.py
@@ -39,14 +39,14 @@ def test_abs():
     @executable
     def abs_tuple(a: jax.Array) -> tuple[jax.Array]: ...
 
-    assert abs_tuple(array(-2, dtype=jax.numpy.int32))[0] == array(
-        2, dtype=jax.numpy.int32
+    assert abs_tuple(array(-2, dtype=jax.numpy.int32)) == (
+        array(2, dtype=jax.numpy.int32),
     )
-    assert abs_tuple(array(0, dtype=jax.numpy.int32))[0] == array(
-        0, dtype=jax.numpy.int32
+    assert abs_tuple(array(0, dtype=jax.numpy.int32)) == (
+        array(0, dtype=jax.numpy.int32),
     )
-    assert abs_tuple(array(2, dtype=jax.numpy.int32))[0] == array(
-        2, dtype=jax.numpy.int32
+    assert abs_tuple(array(2, dtype=jax.numpy.int32)) == (
+        array(2, dtype=jax.numpy.int32),
     )
 
     @executable