Address review comments by @mtsokol.

Author: hameerabbasi

sparse/mlir_backend/_ops.py

@@ -50,7 +50,7 @@ def add(a, b):
                             for t in (a_tensor_type, b_tensor_type, out_tensor_type)
                         ]
                     ),
-                    ir.ArrayAttr.get([ir.Attribute.parse("#linalg.iterator_type<parallel>")] * out_tensor_type.rank),
+                    ir.ArrayAttr.get([ir.Attribute.parse("#linalg.iterator_type<parallel>")] * max_rank),
                 )
                 block = generic_op.regions[0].blocks.append(dtype, dtype, dtype)
                 with ir.InsertionPoint(block):
@@ -160,11 +160,11 @@ def convert(in_tensor):
 
 
 def add(x1: Array, x2: Array, /) -> Array:
+    # TODO: Determine output format via autoscheduler
     ret_storage_format = _determine_format(x1.format, x2.format, dtype=x1.dtype, union=True)
     ret_storage = ret_storage_format._get_ctypes_type(owns_memory=True)()
     out_tensor_type = ret_storage_format._get_mlir_type(shape=np.broadcast_shapes(x1.shape, x2.shape))
 
-    # TODO: Determine output format via autoscheduler
     add_module = get_add_module(
         x1._get_mlir_type(),
         x2._get_mlir_type(),

sparse/mlir_backend/levels.py

@@ -212,6 +212,7 @@ def _get_storage_format(
 
 
 def _is_sparse_level(lvl: Level | LevelFormat, /) -> bool:
+    assert isinstance(lvl, Level | LevelFormat)
     if isinstance(lvl, Level):
         lvl = lvl.format
     return LevelFormat.Dense != lvl
@@ -226,6 +227,18 @@ def _count_dense_levels(format: StorageFormat) -> int:
 
 
 def _determine_format(*formats: StorageFormat, dtype: DType, union: bool, out_ndim: int | None = None) -> StorageFormat:
+    """Determines the output format from a group of input formats.
+
+    1. Counts the sparse levels for `union=True`, and dense ones for `union=False`.
+    2. Gets the max number of counted levels for each format.
+    3. Constructs a format with the same number of counted levels.
+       Sparse levels are replaced with `LevelFormat.Compressed`.
+
+    Returns
+    -------
+    StorageFormat
+        Output storage format.
+    """
     if len(formats) == 0:
         if out_ndim is None:
             out_ndim = 0