diff --git a/mlir/test/Integration/Dialect/SparseTensor/taco/test_scalar_tensor_algebra.py b/mlir/test/Integration/Dialect/SparseTensor/taco/test_scalar_tensor_algebra.py --- a/mlir/test/Integration/Dialect/SparseTensor/taco/test_scalar_tensor_algebra.py +++ b/mlir/test/Integration/Dialect/SparseTensor/taco/test_scalar_tensor_algebra.py @@ -24,5 +24,13 @@ passed = np.array_equal(indices, [[0, 1], [1, 2]]) passed += np.array_equal(values, [30.0, 120.0]) -# CHECK: Number of passed: 2 +# Sum all the values in A. +S[0] = A[i, j] +passed += (S.get_scalar_value() == 50.0) + +indices, values = S.get_coordinates_and_values() +passed += (len(indices)==0) +passed += (values == 50.0) + +# CHECK: Number of passed: 5 print("Number of passed:", passed) diff --git a/mlir/test/Integration/Dialect/SparseTensor/taco/tools/mlir_pytaco.py b/mlir/test/Integration/Dialect/SparseTensor/taco/tools/mlir_pytaco.py --- a/mlir/test/Integration/Dialect/SparseTensor/taco/tools/mlir_pytaco.py +++ b/mlir/test/Integration/Dialect/SparseTensor/taco/tools/mlir_pytaco.py @@ -477,7 +477,7 @@ def emit_tensor_init(self) -> ir.RankedTensorType: """Returns an initialization for the destination tensor.""" - if self.dst_format is None: + if self.dst_format is None or self.dst_format.rank() == 0: # Initialize the dense tensor. ir_type = _mlir_type_from_taco_type(self.dst_dtype) tensor = linalg.InitTensorOp(self.dst_dims, ir_type).result @@ -1023,11 +1023,26 @@ return ctypes.pointer(ctypes.cast(ptr, ctypes.c_void_p)) + def get_scalar_value(self) -> _AnyRuntimeType: + """Returns the value for the scalar tensor. + + This method also evaluates the assignment to the tensor. + + Raises: + ValueError: If the tensor is not a scalar. + """ + if self.order != 0: + raise ValueError(f"Expected a scalar tensor, got: rank={self.order}") + + self._sync_value() + return self._dense_storage + + def get_coordinates_and_values( self) -> Tuple[List[Tuple[int, ...]], List[_AnyRuntimeType]]: """Returns the coordinates and values for the non-zero elements. - This method also evaluate the assignment to the tensor and unpack the + This method also evaluates the assignment to the tensor and unpack the sparse tensor. """ self._sync_value() @@ -1036,6 +1051,9 @@ self.unpack() return (self._coords, self._values) + if self.order == 0: + return ([], self._dense_storage) + # Coordinates for non-zero elements, grouped by dimensions. coords_by_dims = self._dense_storage.nonzero() # Coordinates for non-zero elements, grouped by elements.