Fixing tests

patrick-kidger · Jul 10, 2023 · bed8e98 · bed8e98
1 parent 77f0b76
commit bed8e98
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 5 deletions.
diff --git a/tests/requirements.txt b/tests/requirements.txt
@@ -1,5 +1,3 @@
 jaxlib
 pytest
 beartype
-tensorflow
-tf2onnx
diff --git a/tests/test_onnx.py b/tests/test_onnx.py
@@ -1,9 +1,12 @@
 import jax
 import jax.numpy as jnp
+import pytest
 
 import equinox.internal as eqxi
 
 
+# Currently failing due to a bug in tf2onnx
+@pytest.mark.skip
 def test_onnx_export():
     @jax.vmap
     def fn(x, y):

diff --git a/tests/test_while_loop.py b/tests/test_while_loop.py
@@ -702,9 +702,9 @@ def run(init_carry):
     run((1.0, 2.0, jnp.array(3.0), jnp.array(4.0)))
     text, _ = capfd.readouterr()
     assert (
-        "symbolic_zero_gradient (True, True, (False, True, False, True, True, True))"
-        in text
-    )
+        "symbolic_zero_gradient "
+        "(True, True, True, (False, True, False, True, True, True))"
+    ) in text
 
 
 def test_buffer_index():