[mlir][sparse][taco] Support true dense tensors and all dense sparse tensors.

[bixia]

Summary

The only method to create a true dense tensor (i.e un-annotated) in MLIR-PyTACO
is through the from_array method. However, the annotated all dense tensors are
also implemented as true dense tensor currently. The PR fixes the
implementation to support annotated all dense sparse tensors.

Extend the tensor init method to support the construction of a tensor without
any sparsity annotation.

Change the tensor to_file method to only support writing unpacked sparse
tensors to file through the MLIR sparse tensor dialect.

Add unit tests for true dense tensors and all dense sparse tensors.

[aartbik]

I found the summary a bit confusion at first reading, because some information appears later in the text.
For historical accuracy I would rewrite this a bit, as in

Currently the only method to create a true dense tensor (i.e un-annotated) in MLIR-PyTACO is through
the from_array method. All other dense tensors in PyTACO are implemented as tensor with all dense
dimensions. This PR changes .....

Also, I am pondering a bit if we want to keep the ability to define tensors with all dense dimensions as well,
just like we can at MLIR level. Probably not important for this revision, so we can chat a bit more about it later.

mlir/test/Integration/Dialect/SparseTensor/taco/test_simple_tensor_algebra.py

33

Our tests often give a lot of output with e.g. date and times, so just checking for "2" is very brittle.
I would make the output more precise e.g.

print('Number of passed:', passed)

and check for that

[aartbik]

Per off line explanation, None vs all-dense. Please add a test case test_Dense.py that shows all different ways explicitly

[bixia]

Address the review comment by: revising the summary; printing more information for FileCheck.

[bixia]

I rewrote the commit message. PTAL.

[aartbik]

Can you please also a pytaco test for the new feature? Both to test, but also to document in a way.

[bixia]

Add a PyTACO test to show the use of true dense tensors.

[bixia]

Fix the comment in a test.

[bixia]

Fix a comment.

[aartbik]

mlir/test/Integration/Dialect/SparseTensor/taco/test_true_dense_tensor_algebra.py

12–13

These two lines are not used, right?

[bixia]

Address review comments: remove two unused lines, add a comment for is_dense=True.

[aartbik]

Ship it!

[bixia]

Rebase.

mlir/test/Integration/Dialect/SparseTensor/taco/test_simple_tensor_algebra.py

 # RUN: SUPPORTLIB=%mlir_runner_utils_dir/libmlir_c_runner_utils%shlibext %PYTHON %s | FileCheck %s
 
+import numpy as np
 import os
 import sys
 
 _SCRIPT_PATH = os.path.dirname(os.path.abspath(__file__))
 sys.path.append(_SCRIPT_PATH)
 from tools import mlir_pytaco_api as pt
 
 compressed = pt.compressed
 dense = pt.dense
 
 # Ensure that we can run an unmodified PyTACO program with a simple tensor
 # algebra expression using tensor index notation, and produce the expected
 # result.
 i, j = pt.get_index_vars(2)
 A = pt.tensor([2, 3])
 B = pt.tensor([2, 3])
 C = pt.tensor([2, 3])
-D = pt.tensor([2, 3], dense)
+D = pt.tensor([2, 3], compressed)
 A.insert([0, 1], 10)
 A.insert([1, 2], 40)
 B.insert([0, 0], 20)
 B.insert([1, 2], 30)
 C.insert([0, 1], 5)
 C.insert([1, 2], 7)
 D[i, j] = A[i, j] + B[i, j] - C[i, j]
 
-# CHECK: [20. 5. 0. 0. 0. 63.]
-print(D.to_array().reshape(6))
+indices, values = D.get_coordinates_and_values()
+passed = np.allclose(indices, [[0, 0], [0, 1], [1, 2]])
+passed += np.allclose(values, [20.0, 5.0, 63.0])
+
+# CHECK: Number of passed: 2
+print("Number of passed:", passed)

mlir/test/Integration/Dialect/SparseTensor/taco/test_true_dense_tensor_algebra.py

+# RUN: SUPPORTLIB=%mlir_runner_utils_dir/libmlir_c_runner_utils%shlibext %PYTHON %s | FileCheck %s
+
+import numpy as np
+import os
+import sys
+
+_SCRIPT_PATH = os.path.dirname(os.path.abspath(__file__))
+sys.path.append(_SCRIPT_PATH)
+from tools import mlir_pytaco_api as pt
+
+i, j = pt.get_index_vars(2)
+# Both tensors are true dense tensors.
+A = pt.from_array(np.full([2,3], 1, dtype=np.float64))
+B = pt.from_array(np.full([2,3], 2, dtype=np.float64))
+# Define the result tensor as a true dense tensor. The parameter is_dense=True
+# is an MLIR-PyTACO extension.
+C = pt.tensor([2, 3], is_dense=True)
+
+C[i, j] = A[i, j] + B[i, j]
+
+# CHECK: [3. 3. 3. 3. 3. 3.]
+print(C.to_array().reshape(6))

mlir/test/Integration/Dialect/SparseTensor/taco/tools/mlir_pytaco.py

def __post_init__(self) -> None:
     if not isinstance(self.ordering, ModeOrdering):
       raise ValueError(f"Expected ModeOrdering: {self.ordering}")
 
     if self.format_pack.rank() != self.ordering.rank():
       raise ValueError("Inconsistent ModeFormatPack and ModeOrdering: "
                        f"len({self.format_pack}) != "
                        f"len({self.ordering})")
 
-  def is_dense(self) -> bool:
-    """Returns true if all the Tensor dimensions have a dense format."""
-    return all([f == ModeFormat.DENSE for f in self.format_pack.formats])
-
   def rank(self) -> int:
     """Returns the number of dimensions represented by the format."""
     return self.format_pack.rank()
 
   def mlir_tensor_attr(self) -> Optional[sparse_tensor.EncodingAttr]:
     """Constructs the MLIR attributes for the tensor format."""
-    if self.is_dense():
-      return None
-
     order = (
         range(self.rank()) if
         (self.ordering is None) else self.ordering.ordering)
     mlir_storage_format = [f.value for f in self.format_pack.formats]
     return sparse_tensor.EncodingAttr.get(mlir_storage_format,
                                           ir.AffineMap.get_permutation(order),
                                           _POINTER_BIT_WIDTH, _INDEX_BIT_WIDTH)
 
class _StructOpInfo:
   Attributes:
     dst_indices: A tuple of IndexVar, representing the result dimensions of the
       structured op. This is used to construct the temporary variable for the
       tensor to hold the structured op result.
     dst_dims: A tuple of int, representing the result shape of the structured
       op.
     dst_dtype: A DType representing the data type of the structured op result.
     dst_name: A string representing the name of the structured op result.
-    dst_format: A Format object representing the destination tensor format.
+    dst_format: An optional Format object representing the destination tensor
+      format. None represents a true dense tensor.
   """
   dst_indices: Tuple[IndexVar, ...]
   dst_dims: Tuple[int, ...]
   dst_dtype: DType
   dst_name: str
-  dst_format: Format
+  dst_format: Optional[Format]
 
   def __post_init__(self) -> None:
     """Verifies the integrity of the attribute values."""
     assert len(self.dst_indices) == len(self.dst_dims)
-    assert self.dst_format is not None
 
   def emit_tensor_init(self) -> ir.RankedTensorType:
     """Returns an initialization for the destination tensor."""
-    if self.dst_format.is_dense():
+    if self.dst_format is None:
       # Initialize the dense tensor.
       ir_type = _mlir_type_from_taco_type(self.dst_dtype)
       tensor = linalg.InitTensorOp(self.dst_dims, ir_type).result
       zero = arith.ConstantOp(ir_type, 0.0)
       return linalg.FillOp(output=tensor, value=zero).results[0]
 
     # Initialize the sparse tensor.
     mlir_type = _mlir_tensor_type(self.dst_dtype, self.dst_dims,
else:
       raise ValueError(f"Invalid format argument: {fmt}.")
 
   def __init__(self,
                value_or_shape: Optional[Union[List[int], Tuple[int, ...], float,
                                               int]] = None,
                fmt: Optional[Union[ModeFormat, List[ModeFormat],
                                    Format]] = None,
                dtype: Optional[DType] = None,
-               name: Optional[str] = None):
+               name: Optional[str] = None,
+               is_dense: bool = False):
     """The tensor constructor interface defined by TACO API.
 
     Args:
       value_or_shape: This argument is optional and can be int, float,
         List[int], or Tuple[int, ...]. If this argument is an int or float,
         creates a scalar tensor and initializes it with the value. If this
         argument is a list or tuple of int, uses it as the shape to create a
         tensor.
       fmt: This argument can be a ModeFormat, List[ModeFormat], or format. If
         this argument is a ModeFormat, uses this ModeFormat for all the tensor
         dimensions. If this argument is a list of ModeFormat, the len of the
         list should equal to the rank of the tensor. If this argument is a
         format, uses it for the format of the tensor.
       dtype: An object of dtype, representing the data type of the tensor.
       name: A string name of the tensor. If a name is not given, creates a
         unique name for the tensor.
+      is_dense: A boolean variable to indicate whether the tensor is a dense
+        tensor without any sparsity annotation.
 
     Raises:
       ValueError: If there is any inconsistency among the input arguments.
     """
-    # Take care of the argument default values.
-    fmt = fmt or ModeFormat.COMPRESSED
+    # Take care of the argument default values common to both sparse tensors
+    # and dense tensors.
     dtype = dtype or DType(Type.FLOAT64)
     self._name = name or self._get_unique_name()
-
-    self._dtype = dtype
     self._assignment = None
+    self._sparse_value_location = _SparseValueInfo._UNPACKED
+    self._dense_storage = None
+    self._dtype = dtype
+
+    if is_dense:
+      assert (fmt is None)
+      assert (isinstance(value_or_shape, tuple) or isinstance(
+          value_or_shape, list)) and _all_instance_of(value_or_shape, int)
+      self._shape = value_or_shape
+      self._format = None
+      return
+
+    fmt = fmt or ModeFormat.COMPRESSED
     # We currently use _coords and _values to host the sparse tensor value with
-    # COO format, and _dense_storage to host the dense tensor value. We haven't
-    # implement the conversion between the two storages yet. This will be
-    # improved in a follow up CL.
+    # COO format, and _dense_storage to host the dense tensor value. We don't
+    # support the conversion between the two storages.
     self._coords = []
     self._values = []
-    self._sparse_value_location = _SparseValueInfo._UNPACKED
-    self._dense_storage = None
     self._stats = _Stats()
     if value_or_shape is None or isinstance(value_or_shape, int) or isinstance(
         value_or_shape, float):
       # Create a scalar tensor and ignore the fmt parameter.
       self._shape = []
       self._format = _make_format([], [])
       if value_or_shape is not None:
         self._dense_storage = np.array(value_or_shape, dtype=self._dtype.value)
def unpack(self) -> None:
     assert np.allclose(self.shape, shape)
     assert nse == len(values)
     self._coords = indices
     self._values = values
     self._sparse_value_location = _SparseValueInfo._UNPACKED
 
   def __repr__(self) -> str:
     self._sync_value()
-    self._unpack()
+    self.unpack()
     value_str = (f"{repr(self._dense_storage)})" if self.is_dense() else
                  f"{repr(self._coords)} {repr(self._values)})")
     return (f"Tensor(_name={repr(self._name)} "
             f"_dtype={repr(self._dtype)} : ") + value_str
 
   def insert(self, coords: List[int], val: Union[float, int]) -> None:
     """Inserts a value to the given coordinate.
 
def insert(self, coords: List[int], val: Union[float, int]) -> None:
 
     if not isinstance(val, int) and not isinstance(val, float):
       raise ValueError(f"Value is neither int nor float: {val}.")
 
     self._coords.append(tuple(coords))
     self._values.append(self._dtype.value(val))
 
   def is_dense(self) -> bool:
-    """Returns true if all the Tensor dimensions have a dense format."""
-    return self._format.is_dense()
+    """Returns true if the tensor doesn't have sparsity annotation."""
+    return self._format is None
 
   def to_array(self) -> np.ndarray:
     """Returns the numpy array for the Tensor.
 
     This is currenly only implemented for dense Tensor.
     """
     if not self.is_dense():
       raise ValueError("Conversion from non-dense Tensor "
def from_array(array: np.ndarray) -> "Tensor":
     Returns:
       A Tensor object.
 
     Raises:
       ValueError if the data type of the numpy array is not float64.
     """
     if array.dtype != np.float64:
       raise ValueError(f"Expected float64 value type: {array.dtype}.")
-    tensor = Tensor(array.shape, ModeFormat.DENSE)
+    tensor = Tensor(array.shape, is_dense=True)
     tensor._dense_storage = np.copy(array)
     return tensor
 
   @staticmethod
   def from_coo(
       coordinates: List[Tuple[int, ...]],
       values: List[_AnyRuntimeType],
       fmt: Format,
class Tensor:
   def to_file(self, filename: str) -> None:
     """Output the tensor value to a file.
 
     This method evaluates any pending assignment to the tensor and outputs the
     tensor value.
 
     Args:
       filename: A string file name.
+
+    Raises:
+       ValueError: If the tensor is dense, or an unpacked sparse tensor.
     """
     self._sync_value()
-    if not self.is_unpacked():
-      utils.output_sparse_tensor(self._packed_sparse_value, filename,
-                                 self._format.format_pack.formats)
-      return
-
-    # TODO: Use MLIR code to output the value.
-    coords, values = self.get_coordinates_and_values()
-    assert len(coords) == len(values)
-    with open(filename, "w") as file:
-      # Output a comment line and the meta data.
-      file.write("; extended FROSTT format\n")
-      file.write(f"{self.order} {len(coords)}\n")
-      file.write(f"{' '.join(map(lambda i: str(i), self.shape))}\n")
-      # Output each (coordinate value) pair in a line.
-      for c, v in zip(coords, values):
-        # The coordinates are 1-based in the text file and 0-based in memory.
-        plus_one_to_str = lambda x: str(x + 1)
-        file.write(f"{' '.join(map(plus_one_to_str,c))} {v}\n")
+
+    if self.is_dense():
+      raise ValueError("Writing dense tensors without sparsity annotation to "
+                       "file is not supported.")
+
+    if self.is_unpacked():
+      raise ValueError("Writing unpacked sparse tensors to file is not "
+                       "supported.")
+
 
   @property
   def dtype(self) -> DType:
     """Returns the data type for the Tensor."""
     return self._dtype
 
   @property
   def format(self) -> Format:
class Tensor:
 
   def _sync_value(self) -> None:
     """Updates the tensor value by evaluating the pending assignment."""
     if self._assignment is not None:
       self.evaluate()
 
   def mlir_tensor_type(self) -> ir.RankedTensorType:
     """Returns the MLIR type for the tensor."""
-    return _mlir_tensor_type(self._dtype, tuple(self._shape),
-                             self._format.mlir_tensor_attr())
+    mlir_attr = None if (
+        self._format is None) else self._format.mlir_tensor_attr()
+    return _mlir_tensor_type(self._dtype, tuple(self._shape), mlir_attr)
 
   def dense_dst_ctype_pointer(self) -> ctypes.pointer:
     """Returns the ctypes pointer for the pointer to an MemRefDescriptor.
 
     For a dense tensor output, the MLIR compiler allocates the storage for
     the tensor. This routine returns the pointer to an MLIR MemRefDescriptor for
     receiving the tensor.
     """
if self.is_unpacked():
       ptr = utils.coo_tensor_to_sparse_tensor(shape, values, indices)
     else:
       ptr = self._packed_sparse_value
 
     return ctypes.pointer(ctypes.cast(ptr, ctypes.c_void_p))
 
   def get_coordinates_and_values(
       self) -> Tuple[List[Tuple[int, ...]], List[_AnyRuntimeType]]:
-    """Returns the coordinates and values for the non-zero elements."""
+    """Returns the coordinates and values for the non-zero elements.
+
+    This method also evaluate the assignment to the tensor and unpack the
+    sparse tensor.
+    """
+    self._sync_value()
+
     if not self.is_dense():
-      assert (self.is_unpacked())
+      self.unpack()
       return (self._coords, self._values)
 
     # Coordinates for non-zero elements, grouped by dimensions.
     coords_by_dims = self._dense_storage.nonzero()
     # Coordinates for non-zero elements, grouped by elements.
     coords = np.transpose(coords_by_dims)
     values = self._dense_storage[coords_by_dims]
     return (coords, values)
) -> None:
   # processing Access objects multiple times by skipping the processing if expr
   # is already in the dictionary.
   if expr in expr_to_info:
     return
 
   if isinstance(expr, Access):
     src_indices = expr.indices
     src_dims = tuple(expr.tensor.shape)
+    if expr.tensor.format is None:
+      # Treat each dimension of a dense tensor as DENSE for the purpose of
+      # calculating temporary tensor storage format.
+      mode_formats = tuple([ModeFormat.DENSE] * len(src_dims))
+    else:
       mode_formats = tuple(expr.tensor.format.format_pack.formats)
     assert len(src_dims) == len(mode_formats)
     dim_infos = tuple([_DimInfo(d, m) for d, m in zip(src_dims, mode_formats)])
   else:
     assert isinstance(expr, _BinaryExpr)
     a_info = expr_to_info[expr.a]
     index_to_dim_info = {
         i: d for i, d in zip(a_info.src_indices, a_info.dim_infos)
     }

mlir/test/Integration/Dialect/SparseTensor/taco/tools/mlir_pytaco_io.py

Raises:
     instance of Format or fmt is not a sparse tensor.
   """
   if (not isinstance(filename, str) or
       (not filename.endswith(_MTX_FILENAME_SUFFIX) and
        not filename.endswith(_TNS_FILENAME_SUFFIX))):
     raise ValueError("Expected string filename ends with "
                      f"{_MTX_FILENAME_SUFFIX} or {_TNS_FILENAME_SUFFIX}: "
                      f"{filename}.")
-  if not isinstance(fmt, Format) or fmt.is_dense():
-    raise ValueError(f"Expected a sparse Format object: {fmt}.")
 
   return Tensor.from_file(filename, fmt, DType(Type.FLOAT64))
 
 
 def write(filename: str, tensor: Tensor) -> None:
   """Outputs a tensor to a given file.
 
   The name suffix of the file specifies the format of the output. We currently

mlir/test/Integration/Dialect/SparseTensor/taco/tools/testing_utils.py

def compare_sparse_tns(expected: str, actual: str, rtol: float = 0.0001) -> bool:
   expected_data = np.loadtxt(expected, np.float64, skiprows=3)
   return np.allclose(actual_data, expected_data, rtol=rtol)
 
 
 def file_as_string(file: str) -> str:
   """Returns contents of file as string."""
   with open(file, "r") as f:
     return f.read()
+
+
+def run_test(f):
+  """Prints the test name and runs the test."""
+  print(f.__name__)
+  f()
+  return f

mlir/test/Integration/Dialect/SparseTensor/taco/unit_test_tensor_core.py

+# RUN: SUPPORTLIB=%mlir_runner_utils_dir/libmlir_c_runner_utils%shlibext %PYTHON %s | FileCheck %s
+
+from string import Template
+
+import numpy as np
+import os
+import sys
+import tempfile
+
+_SCRIPT_PATH = os.path.dirname(os.path.abspath(__file__))
+sys.path.append(_SCRIPT_PATH)
+from tools import mlir_pytaco
+from tools import testing_utils as testing_utils
+
+# Define the aliases to shorten the code.
+_COMPRESSED = mlir_pytaco.ModeFormat.COMPRESSED
+_DENSE = mlir_pytaco.ModeFormat.DENSE
+
+
+# CHECK-LABEL: test_tensor_all_dense_sparse
+@testing_utils.run_test
+def test_tensor_all_dense_sparse():
+  a = mlir_pytaco.Tensor([4], [_DENSE])
+  passed = (not a.is_dense())
+  passed += (a.order == 1)
+  passed += (a.shape[0] == 4)
+  # CHECK: Number of passed: 3
+  print("Number of passed:", passed)
+
+
+# CHECK-LABEL: test_tensor_true_dense
+@testing_utils.run_test
+def test_tensor_true_dense():
+  a = mlir_pytaco.Tensor.from_array(np.random.uniform(size=5))
+  passed = a.is_dense()
+  passed += (a.order == 1)
+  passed += (a.shape[0] == 5)
+  # CHECK: Number of passed: 3
+  print("Number of passed:", passed)

mlir/test/Integration/Dialect/SparseTensor/taco/unit_test_tensor_io.py

 # RUN: SUPPORTLIB=%mlir_runner_utils_dir/libmlir_c_runner_utils%shlibext %PYTHON %s | FileCheck %s
 
 from string import Template
 
 import numpy as np
 import os
 import sys
 import tempfile
 
 _SCRIPT_PATH = os.path.dirname(os.path.abspath(__file__))
 sys.path.append(_SCRIPT_PATH)
 from tools import mlir_pytaco
 from tools import mlir_pytaco_io
 from tools import mlir_pytaco_utils as pytaco_utils
+from tools import testing_utils as testing_utils
+
 
 # Define the aliases to shorten the code.
 _COMPRESSED = mlir_pytaco.ModeFormat.COMPRESSED
 _DENSE = mlir_pytaco.ModeFormat.DENSE
 
 
-def _run(f):
-  print(f.__name__)
-  f()
-  return f
-
-
 _FORMAT = mlir_pytaco.Format([_COMPRESSED, _COMPRESSED])
 _MTX_DATA_TEMPLATE = Template(
     """%%MatrixMarket matrix coordinate real $general_or_symmetry
 3 3 3
 3 1 3
 1 2 2
 3 2 4
 """)
 
 
 def _get_mtx_data(value):
   mtx_data = _MTX_DATA_TEMPLATE
   return mtx_data.substitute(general_or_symmetry=value)
 
 
 # CHECK-LABEL: test_read_mtx_matrix_general
-@_run
+@testing_utils.run_test
 def test_read_mtx_matrix_general():
   with tempfile.TemporaryDirectory() as test_dir:
     file_name = os.path.join(test_dir, "data.mtx")
     with open(file_name, "w") as file:
       file.write(_get_mtx_data("general"))
     a = mlir_pytaco_io.read(file_name, _FORMAT)
   passed = 0
   # The value of a is stored as an MLIR sparse tensor.
   passed += (not a.is_unpacked())
   a.unpack()
   passed += (a.is_unpacked())
   coords, values = a.get_coordinates_and_values()
   passed += np.allclose(coords, [[0, 1], [2, 0], [2, 1]])
   passed += np.allclose(values, [2.0, 3.0, 4.0])
   # CHECK: 4
   print(passed)
 
 
 # CHECK-LABEL: test_read_mtx_matrix_symmetry
-@_run
+@testing_utils.run_test
 def test_read_mtx_matrix_symmetry():
   with tempfile.TemporaryDirectory() as test_dir:
     file_name = os.path.join(test_dir, "data.mtx")
     with open(file_name, "w") as file:
       file.write(_get_mtx_data("symmetric"))
     a = mlir_pytaco_io.read(file_name, _FORMAT)
   passed = 0
   # The value of a is stored as an MLIR sparse tensor.

 3 2
 3 1 3
 1 2 2
 3 2 4
 """
 
 
 # CHECK-LABEL: test_read_tns
-@_run
+@testing_utils.run_test
 def test_read_tns():
   with tempfile.TemporaryDirectory() as test_dir:
     file_name = os.path.join(test_dir, "data.tns")
     with open(file_name, "w") as file:
       file.write(_TNS_DATA)
     a = mlir_pytaco_io.read(file_name, _FORMAT)
   passed = 0
   # The value of a is stored as an MLIR sparse tensor.
   passed += (not a.is_unpacked())
   a.unpack()
   passed += (a.is_unpacked())
   coords, values = a.get_coordinates_and_values()
   passed += np.allclose(coords, [[0, 1], [2, 0], [2, 1]])
   passed += np.allclose(values, [2.0, 3.0, 4.0])
   # CHECK: 4
   print(passed)
 
 
 # CHECK-LABEL: test_write_unpacked_tns
-@_run
+@testing_utils.run_test
 def test_write_unpacked_tns():
   a = mlir_pytaco.Tensor([2, 3])
   a.insert([0, 1], 10)
   a.insert([1, 2], 40)
   a.insert([0, 0], 20)
   with tempfile.TemporaryDirectory() as test_dir:
     file_name = os.path.join(test_dir, "data.tns")
+    try:
       mlir_pytaco_io.write(file_name, a)
-    with open(file_name, "r") as file:
-      lines = file.readlines()
-  passed = 0
-  # Skip the comment line in the output.
-  if lines[1:] == ["2 3\n", "2 3\n", "1 2 10.0\n", "2 3 40.0\n", "1 1 20.0\n"]:
-    passed = 1
-  # CHECK: 1
-  print(passed)
+    except ValueError as e:
+      # CHECK: Writing unpacked sparse tensors to file is not supported
+      print(e)
 
 
 # CHECK-LABEL: test_write_packed_tns
-@_run
+@testing_utils.run_test
 def test_write_packed_tns():
   a = mlir_pytaco.Tensor([2, 3])
   a.insert([0, 1], 10)
   a.insert([1, 2], 40)
   a.insert([0, 0], 20)
   b = mlir_pytaco.Tensor([2, 3])
   i, j = mlir_pytaco.get_index_vars(2)
   b[i, j] = a[i, j] + a[i, j]