Merge branch 'sasha/jax/compile-init' into sasha/jax/jit

xdslproject · Aug 17, 2024 · 26bc15e · 26bc15e
2 parents 18f7fe2 + d8555c2
commit 26bc15e
Showing 1 changed file with 9 additions and 9 deletions.
diff --git a/tests/backend/test_jax_executable.py b/tests/backend/test_jax_executable.py
@@ -26,15 +26,15 @@ def test_abs():
 
     executable = JaxExecutable.compile(module)
 
-    assert executable.execute([array(-2, dtype=jax.numpy.int32)])[0] == array(
-        2, dtype=jax.numpy.int32
-    )
-    assert executable.execute([array(0, dtype=jax.numpy.int32)])[0] == array(
-        0, dtype=jax.numpy.int32
-    )
-    assert executable.execute([array(2, dtype=jax.numpy.int32)])[0] == array(
-        2, dtype=jax.numpy.int32
-    )
+    assert executable.execute([array(-2, dtype=jax.numpy.int32)]) == [
+        array(2, dtype=jax.numpy.int32)
+    ]
+    assert executable.execute([array(0, dtype=jax.numpy.int32)]) == [
+        array(0, dtype=jax.numpy.int32)
+    ]
+    assert executable.execute([array(2, dtype=jax.numpy.int32)]) == [
+        array(2, dtype=jax.numpy.int32)
+    ]
 
     @executable
     def abs_tuple(a: jax.Array) -> tuple[jax.Array]: ...