[Relax][Frontend][TFLite] Fix and test MATRIX_DIAG, MATRIX_SET_DIAG, SPARSE_TO_DENSE (#19408)

This PR partially implements test coverage requested in issue #18971 for
Relax TFLite frontend operator tests.

## Bug Fix

The TFLite frontend converters for `MATRIX_DIAG`, `MATRIX_SET_DIAG`, and
`SPARSE_TO_DENSE` were broken due to calling non-existent Relax ops:

- `relax.op.matrix_set_diag` - never registered in Relax
- `relax.op.sparse_to_dense` - never registered in Relax

These ops only exist as TOPI packed functions (`topi.matrix_set_diag`,
`topi.sparse_to_dense`).

**Fix:** Replace direct op calls with `call_dps_packed` to invoke the
TOPI packed functions:
- `convert_matrix_diag`: zeros + call_dps_packed("topi.matrix_set_diag",
...)
- `convert_matrix_set_diag`: call_dps_packed("topi.matrix_set_diag",
...)
- `convert_sparse_to_dense`: call_dps_packed("topi.sparse_to_dense",
...)

Refs: #18971

Author: Aharrypotter

python/tvm/relax/frontend/tflite/tflite_frontend.py

@@ -3015,11 +3015,20 @@ def convert_sparse_to_dense(self, op):
             t_type = t.tensor.Type()
             assert t_type in (TensorType.INT32, TensorType.INT64)
 
-        out = relax.op.sparse_to_dense(
-            self.get_tensor_expr(indices),
-            list(self.get_tensor_value(output_shape)),
-            self.get_tensor_expr(values),
-            self.get_tensor_expr(default_value),
+        output_tensors = self.get_output_tensors(op)
+        output_tensor = output_tensors[0]
+        output_shape_val = to_int_list(self.get_tensor_shape(output_tensor))
+        output_dtype = self.get_tensor_type_str(output_tensor.tensor.Type())
+
+        indices_expr = self.get_tensor_expr(indices)
+        values_expr = self.get_tensor_expr(values)
+        default_value_expr = self.get_tensor_expr(default_value)
+        output_shape_expr = relax.const(list(self.get_tensor_value(output_shape)), "int32")
+
+        out = relax.op.call_dps_packed(
+            "topi.sparse_to_dense",
+            (indices_expr, output_shape_expr, values_expr, default_value_expr),
+            out_sinfo=relax.TensorStructInfo(output_shape_val, output_dtype),
         )
 
         return out
@@ -3700,7 +3709,18 @@ def convert_matrix_set_diag(self, op):
         input_expr = self.get_tensor_expr(input_tensors[0])
         diagonal_expr = self.get_tensor_expr(input_tensors[1])
 
-        out = relax.op.matrix_set_diag(input_expr, diagonal_expr)
+        output_tensors = self.get_output_tensors(op)
+        output_tensor = output_tensors[0]
+        output_shape = to_int_list(self.get_tensor_shape(output_tensor))
+        output_dtype = self.get_tensor_type_str(output_tensor.tensor.Type())
+
+        # topi.matrix_set_diag(input, diagonal, k1, k2, super_diag_right_align, sub_diag_right_align)
+        # TFLite MATRIX_SET_DIAG only sets the main diagonal, so k1=0, k2=0
+        out = relax.op.call_dps_packed(
+            "topi.matrix_set_diag",
+            (input_expr, diagonal_expr, relax.const(0), relax.const(0), relax.const(False), relax.const(False)),
+            out_sinfo=relax.TensorStructInfo(output_shape, output_dtype),
+        )
         return out
 
     def convert_matrix_diag(self, op):
@@ -3718,14 +3738,21 @@ def convert_matrix_diag(self, op):
                     scale and zero points to be equal"
             )
 
-        shape = to_int_list(self.get_tensor_shape(diagonal))
-        shape = np.append(shape, shape[-1])
-        dtype = self.get_tensor_type_str(diagonal.tensor.Type())
+        output_tensors = self.get_output_tensors(op)
+        output_tensor = output_tensors[0]
+        output_shape = to_int_list(self.get_tensor_shape(output_tensor))
+        output_dtype = self.get_tensor_type_str(output_tensor.tensor.Type())
 
-        input_expr = relax.op.zeros(tuple(shape), dtype)
         diagonal_expr = self.get_tensor_expr(diagonal)
-
-        out = relax.op.matrix_set_diag(input_expr, diagonal_expr)
+        zeros_expr = relax.op.zeros(output_shape, output_dtype)
+
+        # topi.matrix_set_diag(input, diagonal, k1, k2, super_diag_right_align, sub_diag_right_align)
+        # TFLite MATRIX_DIAG only sets the main diagonal, so k1=0, k2=0
+        out = relax.op.call_dps_packed(
+            "topi.matrix_set_diag",
+            (zeros_expr, diagonal_expr, relax.const(0), relax.const(0), relax.const(False), relax.const(False)),
+            out_sinfo=relax.TensorStructInfo(output_shape, output_dtype),
+        )
         return out
 
     def convert_densify(self, op):

tests/python/relax/test_frontend_tflite.py

@@ -2292,5 +2292,119 @@ def main(x: R.Tensor((1, 4, 4, 3), dtype="float32")) -> R.Tensor(
     tvm.ir.assert_structural_equal(mod, Expected)
 
 
+def test_matrix_diag():
+    """Test TFLite MATRIX_DIAG operator."""
+
+    class MatrixDiag(tf.Module):
+        @tf.function(input_signature=[tf.TensorSpec(shape=(3,), dtype=tf.float32)])
+        def func(self, diagonal):
+            return tf.raw_ops.MatrixDiag(diagonal=diagonal)
+
+    @I.ir_module
+    class Expected:
+        @R.function
+        def main(diagonal: R.Tensor((3,), dtype="float32")) -> R.Tensor((3, 3), dtype="float32"):
+            R.func_attr({"num_input": 1})
+            with R.dataflow():
+                lv: R.Tensor((3, 3), dtype="float32") = R.zeros(R.shape([3, 3]), dtype="float32")
+                gv = R.call_dps_packed(
+                    "topi.matrix_set_diag",
+                    (
+                        lv,
+                        diagonal,
+                        R.const(0, "int32"),
+                        R.const(0, "int32"),
+                        R.const(False, "bool"),
+                        R.const(False, "bool"),
+                    ),
+                    out_sinfo=R.Tensor((3, 3), dtype="float32"),
+                )
+                R.output(gv)
+            return gv
+
+    verify(MatrixDiag, Expected)
+
+
+def test_matrix_set_diag():
+    """Test TFLite MATRIX_SET_DIAG operator."""
+
+    class MatrixSetDiag(tf.Module):
+        @tf.function(
+            input_signature=[
+                tf.TensorSpec(shape=(3, 3), dtype=tf.float32),
+                tf.TensorSpec(shape=(3,), dtype=tf.float32),
+            ]
+        )
+        def func(self, input, diagonal):
+            return tf.raw_ops.MatrixSetDiag(input=input, diagonal=diagonal)
+
+    @I.ir_module
+    class Expected:
+        @R.function
+        def main(
+            input: R.Tensor((3, 3), dtype="float32"),
+            diagonal: R.Tensor((3,), dtype="float32"),
+        ) -> R.Tensor((3, 3), dtype="float32"):
+            R.func_attr({"num_input": 2})
+            with R.dataflow():
+                gv = R.call_dps_packed(
+                    "topi.matrix_set_diag",
+                    (
+                        input,
+                        diagonal,
+                        R.const(0, "int32"),
+                        R.const(0, "int32"),
+                        R.const(False, "bool"),
+                        R.const(False, "bool"),
+                    ),
+                    out_sinfo=R.Tensor((3, 3), dtype="float32"),
+                )
+                R.output(gv)
+            return gv
+
+    verify(MatrixSetDiag, Expected)
+
+
+def test_sparse_to_dense():
+    """Test TFLite SPARSE_TO_DENSE operator."""
+
+    class SparseToDense(tf.Module):
+        @tf.function(
+            input_signature=[
+                tf.TensorSpec(shape=(2,), dtype=tf.int32),
+                tf.TensorSpec(shape=(2,), dtype=tf.float32),
+                tf.TensorSpec(shape=(), dtype=tf.float32),
+            ]
+        )
+        def func(self, indices, values, default_value):
+            # output_shape is provided as a constant, not an input
+            return tf.raw_ops.SparseToDense(
+                sparse_indices=indices,
+                output_shape=tf.constant([3], dtype=tf.int32),
+                sparse_values=values,
+                default_value=default_value,
+            )
+
+    @I.ir_module
+    class Expected:
+        @R.function
+        def main(
+            indices: R.Tensor((2,), dtype="int32"),
+            values: R.Tensor((2,), dtype="float32"),
+            default_value: R.Tensor((), dtype="float32"),
+        ) -> R.Tensor((3,), dtype="float32"):
+            R.func_attr({"num_input": 3})
+            with R.dataflow():
+                gv = R.call_dps_packed(
+                    "topi.sparse_to_dense",
+                    (indices, R.const([3], "int32"), values, default_value),
+                    out_sinfo=R.Tensor((3,), dtype="float32"),
+                )
+                R.output(gv)
+            return gv
+
+    verify(SparseToDense, Expected)
+
+
 if __name__ == "__main__":
     pytest.main(["-s", __file__])