This is an automated email from the ASF dual-hosted git repository.
ruihangl pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new 75a6b308c6 [Relax][Frontend][TFLite] Fix and test `MATRIX_DIAG`,
`MATRIX_SET_DIAG`, `SPARSE_TO_DENSE` (#19408)
75a6b308c6 is described below
commit 75a6b308c6ff1c7838e88b56fbcc0bd8d9f3328b
Author: HoYi <[email protected]>
AuthorDate: Fri Apr 17 00:13:10 2026 +0800
[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
---
.../tvm/relax/frontend/tflite/tflite_frontend.py | 51 ++++++---
tests/python/relax/test_frontend_tflite.py | 114 +++++++++++++++++++++
2 files changed, 153 insertions(+), 12 deletions(-)
diff --git a/python/tvm/relax/frontend/tflite/tflite_frontend.py
b/python/tvm/relax/frontend/tflite/tflite_frontend.py
index ce74f707cf..22f45e8a5a 100644
--- a/python/tvm/relax/frontend/tflite/tflite_frontend.py
+++ b/python/tvm/relax/frontend/tflite/tflite_frontend.py
@@ -3015,11 +3015,20 @@ class OperatorConverter:
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 @@ class OperatorConverter:
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 @@ class OperatorConverter:
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):
diff --git a/tests/python/relax/test_frontend_tflite.py
b/tests/python/relax/test_frontend_tflite.py
index a116daebb1..de26bae25e 100644
--- a/tests/python/relax/test_frontend_tflite.py
+++ b/tests/python/relax/test_frontend_tflite.py
@@ -2292,5 +2292,119 @@ def test_prelu(shared_axes):
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__])
