Fixed flaky test

patrick-kidger · Dec 18, 2023 · ad76799 · ad76799
1 parent b81a55a
commit ad76799
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/tests/test_sparse.py b/tests/test_sparse.py
@@ -98,8 +98,8 @@ def test_add_sparse(getkey):
     assert z_mat.shape == (3, 2, 5, 1, 4)
     assert w.shape == (2, 5, 1, 4)
     assert w_mat.shape == (2, 5, 1, 4)
-    assert jnp.array_equal(w_mat, x_mat * 2)
-    assert jnp.array_equal(z_mat, x_mat + y_mat)
+    assert jnp.allclose(w_mat, x_mat * 2)
+    assert jnp.allclose(z_mat, x_mat + y_mat)
 
 
 def test_mul(getkey):