fixes

Author: crusaderky

tests/test_at.py

@@ -3,10 +3,18 @@
 import numpy as np
 import pytest
 
-from array_api_compat import array_namespace, at, is_dask_array, is_jax_array, is_writeable_array
+from array_api_compat import array_namespace, at, is_dask_array, is_jax_array, is_pydata_sparse_array, is_writeable_array
 from ._helpers import import_, all_libraries
 
 
+def assert_array_equal(a, b):
+    if is_pydata_sparse_array(a):
+        a = a.todense()
+    elif is_dask_array(a):
+        a = a.compute()
+    np.testing.assert_array_equal(a, b)
+
+
 @contextmanager
 def assert_copy(x, copy: bool | None):
     # dask arrays are writeable, but writing to them will hot-swap the
@@ -21,10 +29,13 @@ def assert_copy(x, copy: bool | None):
     x_orig = xp.asarray(x, copy=True)
     yield
 
-    expect_copy = (
-        copy if copy is not None else (not is_writeable_array(x) or is_dask_array(x))
-    )
-    np.testing.assert_array_equal((x == x_orig).all(), expect_copy)
+    if is_dask_array(x):
+        expect_copy = True
+    elif copy is None:
+        expect_copy = not is_writeable_array(x)
+    else:
+        expect_copy = copy
+    assert_array_equal((x == x_orig).all(), expect_copy)
 
 
 @pytest.fixture(params=all_libraries + ["np_readonly"])
@@ -58,15 +69,15 @@ def test_operations(x, copy, op, arg, expect):
     with assert_copy(x, copy):
         y = getattr(at(x, slice(1, None)), op)(arg, copy=copy)
         assert isinstance(y, type(x))
-        np.testing.assert_equal(y, expect)
+        assert_array_equal(y, expect)
 
 
 @pytest.mark.parametrize("copy", [True, False, None])
 def test_get(x, copy):
     with assert_copy(x, copy):
         y = at(x, slice(2)).get(copy=copy)
         assert isinstance(y, type(x))
-        np.testing.assert_array_equal(y, [10, 20])
+        assert_array_equal(y, [10, 20])
         # Let assert_copy test that y is a view or copy
         with suppress((TypeError, ValueError)):
             y[0] = 40
@@ -97,15 +108,15 @@ def test_get_fancy_indices(x, idx, wrap_index):
     with assert_copy(x, True):
         y = at(x, [0, 1]).get()
         assert isinstance(y, type(x))
-        np.testing.assert_array_equal(y, [10, 20])
+        assert_array_equal(y, [10, 20])
         # Let assert_copy test that y is a view or copy
         with suppress((TypeError, ValueError)):
             y[0] = 40
 
     with assert_copy(x, True):
         y = at(x, [0, 1]).get(copy=None)
         assert isinstance(y, type(x))
-        np.testing.assert_array_equal(y, [10, 20])
+        assert_array_equal(y, [10, 20])
         # Let assert_copy test that y is a view or copy
         with suppress((TypeError, ValueError)):
             y[0] = 40
@@ -119,7 +130,7 @@ def test_variant_index_syntax(x, copy):
     with assert_copy(x, copy):
         y = at(x)[:2].set(40, copy=copy)
         assert isinstance(y, type(x))
-        np.testing.assert_array_equal(y, [40, 40, 30])
+        assert_array_equal(y, [40, 40, 30])
     with pytest.raises(ValueError):
         at(x, 1)[2]
     with pytest.raises(ValueError):