fixup! Avoid JIT compilation errors

nengo · Jul 11, 2024 · 38bc180 · 38bc180
1 parent a233133
commit 38bc180
Showing 1 changed file with 2 additions and 4 deletions.
diff --git a/keras_lmu/tests/test_layers.py b/keras_lmu/tests/test_layers.py
@@ -114,10 +114,8 @@ def hidden_cell():
     ):
         assert np.allclose(w0.numpy(), w1.numpy())
 
-    assert np.allclose(cell_out, lmu_cell(inp))
-    assert np.allclose(cell_out, layer_out, atol=3e-6 if feedforward else 1e-8), np.max(
-        np.abs(cell_out - layer_out)
-    )
+    np.testing.assert_allclose(cell_out, lmu_cell(inp))
+    np.testing.assert_allclose(cell_out, layer_out, atol=3e-6 if feedforward else 1e-8)
 
 
 @pytest.mark.parametrize("discretizer", ("zoh", "euler"))