Better tests and fixes.

Author: hameerabbasi

sparse/mlir_backend/_ops.py

@@ -156,11 +156,11 @@ def add(x1: Array, x2: Array) -> Array:
     return Array(storage=ret_storage, shape=tuple(out_tensor_type.shape))
 
 
-def reshape(x: Array, /, shape: tuple[int, ...]):
+def reshape(x: Array, /, shape: tuple[int, ...]) -> Array:
     from ._conversions import _from_numpy
 
     shape = as_shape(shape)
-    ret_storage_format = _determine_format(x.format, dtype=x.dtype, union=len(shape) >= x.ndim, out_ndim=len(shape))
+    ret_storage_format = _determine_format(x.format, dtype=x.dtype, union=len(shape) > x.ndim, out_ndim=len(shape))
     shape_array = _from_numpy(np.asarray(shape, dtype=np.uint64))
     out_tensor_type = ret_storage_format._get_mlir_type(shape=shape)
     ret_storage = ret_storage_format._get_ctypes_type(owns_memory=True)()

sparse/mlir_backend/levels.py

@@ -242,12 +242,11 @@ def _determine_format(*formats: StorageFormat, dtype: DType, union: bool, out_nd
 
     pos_width = 0
     crd_width = 0
-    op = min if union else max
     counter = _count_sparse_levels if not union else _count_dense_levels
     n_counted = None
     order = ()
     for fmt in formats:
-        n_counted = counter(fmt) if n_counted is None else op(n_counted, counter(fmt))
+        n_counted = counter(fmt) if n_counted is None else max(n_counted, counter(fmt))
         pos_width = max(pos_width, fmt.pos_width)
         crd_width = max(crd_width, fmt.crd_width)
         if order != "C":
@@ -256,12 +255,14 @@ def _determine_format(*formats: StorageFormat, dtype: DType, union: bool, out_nd
             elif order[: len(fmt.order)] != fmt.order:
                 order = "C"
 
-    order = order + tuple(range(len(order), out_ndim))
+    if not isinstance(order, str):
+        order = order + tuple(range(len(order), out_ndim))
+        order = order[:out_ndim]
 
     if out_ndim < n_counted:
         n_counted = out_ndim
 
-    n_sparse = n_counted if union else out_ndim - n_counted
+    n_sparse = n_counted if not union else out_ndim - n_counted
 
     levels = (Level(LevelFormat.Dense),) * (out_ndim - n_sparse) + (Level(LevelFormat.Compressed),) * n_sparse
     return get_storage_format(

sparse/mlir_backend/tests/test_simple.py

@@ -300,37 +300,50 @@ def test_copy():
 
 
 @parametrize_dtypes
-def test_reshape(rng, dtype):
+@pytest.mark.parametrize(
+    "format",
+    [
+        "csr",
+        pytest.param("csc", marks=pytest.mark.xfail(reason="https://github.com/llvm/llvm-project/pull/109135")),
+        pytest.param("coo", marks=pytest.mark.xfail(reason="https://github.com/llvm/llvm-project/pull/109641")),
+    ],
+)
+@pytest.mark.parametrize(
+    ("shape", "new_shape"),
+    [
+        ((100, 50), (25, 200)),
+        ((100, 50), (10, 500, 1)),
+        ((80, 1), (8, 10)),
+        ((80, 1), (80,)),
+    ],
+)
+def test_reshape(rng, dtype, format, shape, new_shape):
     DENSITY = 0.5
     sampler = generate_sampler(dtype, rng)
 
-    # CSR, CSC, COO
-    for shape, new_shape in [
-        ((100, 50), (25, 200)),
-        # ((100, 50), (10, 500, 1)),
-        ((80, 1), (8, 10)),
-        # ((80, 1), (80,)),
-    ]:
-        for format in ["csr", "csc", "coo"]:
-            if format == "coo":
-                # NOTE: Blocked by https://github.com/llvm/llvm-project/pull/109135
-                continue
-            if format == "csc":
-                # NOTE: Blocked by https://github.com/llvm/llvm-project/issues/109641
-                continue
-
-            arr = sps.random_array(
-                shape, density=DENSITY, format=format, dtype=dtype, random_state=rng, data_sampler=sampler
-            )
-            arr.eliminate_zeros()
-            arr.sum_duplicates()
-            tensor = sparse.asarray(arr)
-
-            actual = sparse.to_scipy(sparse.reshape(tensor, shape=new_shape))
-            expected = arr.todense().reshape(new_shape)
-
-            np.testing.assert_array_equal(actual.todense(), expected)
+    arr_sps = sps.random_array(
+        shape, density=DENSITY, format=format, dtype=dtype, random_state=rng, data_sampler=sampler
+    )
+    arr_sps.eliminate_zeros()
+    arr_sps.sum_duplicates()
+    arr = sparse.asarray(arr_sps)
+
+    actual = sparse.reshape(arr, shape=new_shape)
+    assert actual.shape == new_shape
+
+    try:
+        scipy_format = sparse.to_scipy(actual).format
+    except RuntimeError:
+        pytest.xfail("No library to compare to.")
+
+    expected = sparse.asarray(arr_sps.reshape(new_shape).asformat(scipy_format)) if scipy_format is not None else arr
 
+    for x, y in zip(expected.get_constituent_arrays(), actual.get_constituent_arrays(), strict=True):
+        np.testing.assert_array_equal(x, y)
+
+
+@parametrize_dtypes
+def test_reshape_csf(dtype):
     # CSF
     csf_shape = (2, 2, 4)
     csf_format = sparse.levels.get_storage_format(
@@ -372,7 +385,6 @@ def test_reshape(rng, dtype):
         csf_tensor = sparse.from_constituent_arrays(format=csf_format, arrays=arrs, shape=shape)
 
         result = sparse.reshape(csf_tensor, shape=new_shape)
-
         for actual, expected in zip(result.get_constituent_arrays(), expected_arrs, strict=True):
             np.testing.assert_array_equal(actual, expected)