You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by sy...@apache.org on 2025/04/30 11:25:30 UTC
(tvm) branch main updated: Add op support for new_zeros op in Exported Program and fx graph frontend (#17911)
This is an automated email from the ASF dual-hosted git repository.
syfeng 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 0c0cdd8958 Add op support for new_zeros op in Exported Program and fx graph frontend (#17911)
0c0cdd8958 is described below
commit 0c0cdd8958c14fbd472609a32552073653c4bc1b
Author: Pratheesh-04-MCW <pr...@multicorewareinc.com>
AuthorDate: Wed Apr 30 16:55:24 2025 +0530
Add op support for new_zeros op in Exported Program and fx graph frontend (#17911)
---
.../frontend/torch/base_fx_graph_translator.py | 19 ++++++++++++++++
.../frontend/torch/exported_program_translator.py | 1 +
python/tvm/relax/frontend/torch/fx_translator.py | 1 +
.../relax/test_frontend_from_exported_program.py | 23 ++++++++++++++++++++
tests/python/relax/test_frontend_from_fx.py | 25 ++++++++++++++++++++++
5 files changed, 69 insertions(+)
diff --git a/python/tvm/relax/frontend/torch/base_fx_graph_translator.py b/python/tvm/relax/frontend/torch/base_fx_graph_translator.py
index ff5e51da0b..7c9ab43cf9 100644
--- a/python/tvm/relax/frontend/torch/base_fx_graph_translator.py
+++ b/python/tvm/relax/frontend/torch/base_fx_graph_translator.py
@@ -1549,6 +1549,25 @@ class BaseFXGraphImporter(metaclass=abc.ABCMeta):
)
)
+ def _new_zeros(self, node: fx.Node) -> relax.Var:
+ args = self.retrieve_args(node)
+ input_tensor = args[0]
+ size = (
+ args[1]
+ if isinstance(args[1], (list, tuple))
+ else (args[1],)
+ if len(args[1:]) == 1
+ else args[1:]
+ )
+ size = relax.ShapeExpr(size)
+ return self.block_builder.emit(
+ relax.op.full(
+ size,
+ relax.const(0, input_tensor.struct_info.dtype),
+ input_tensor.struct_info.dtype,
+ )
+ )
+
def _ones(self, node: fx.Node) -> relax.Var:
import torch
diff --git a/python/tvm/relax/frontend/torch/exported_program_translator.py b/python/tvm/relax/frontend/torch/exported_program_translator.py
index 7919e288c2..8c2e9aa47a 100644
--- a/python/tvm/relax/frontend/torch/exported_program_translator.py
+++ b/python/tvm/relax/frontend/torch/exported_program_translator.py
@@ -484,6 +484,7 @@ class ExportedProgramImporter(BaseFXGraphImporter):
"masked_fill.Scalar": self._masked_fill,
"masked_fill_.Scalar": self._inplace_masked_fill,
"new_ones.default": self._new_ones,
+ "new_zeros.default": self._new_zeros,
"one_hot.default": self._one_hot,
"ones.default": self._ones,
"ones_like.default": lambda node: self.block_builder.emit(
diff --git a/python/tvm/relax/frontend/torch/fx_translator.py b/python/tvm/relax/frontend/torch/fx_translator.py
index 07b3df20aa..39e562e06a 100644
--- a/python/tvm/relax/frontend/torch/fx_translator.py
+++ b/python/tvm/relax/frontend/torch/fx_translator.py
@@ -829,6 +829,7 @@ class TorchFXImporter(BaseFXGraphImporter):
"masked_fill": self._masked_fill,
"masked_scatter": self._masked_scatter,
"new_ones": self._new_ones,
+ "new_zeros": self._new_zeros,
"ones": self._ones,
"one_hot": self._one_hot,
"ones_like": lambda node: self.block_builder.emit(
diff --git a/tests/python/relax/test_frontend_from_exported_program.py b/tests/python/relax/test_frontend_from_exported_program.py
index b9385d1cc2..77f585606a 100644
--- a/tests/python/relax/test_frontend_from_exported_program.py
+++ b/tests/python/relax/test_frontend_from_exported_program.py
@@ -3775,6 +3775,29 @@ def test_new_ones():
verify_model(NewOnes(), example_args, {}, expected1)
+def test_new_zeros():
+ class NewZeros(torch.nn.Module):
+ def forward(self, x):
+ return x.new_zeros(1, 128, 128)
+
+ @tvm.script.ir_module
+ class expected1:
+ @R.function
+ def main(
+ x: R.Tensor((1, 128, 128), dtype="float32")
+ ) -> R.Tuple(R.Tensor((1, 128, 128), dtype="float32")):
+ with R.dataflow():
+ lv: R.Tensor((1, 128, 128), dtype="float32") = R.full(
+ R.shape([1, 128, 128]), R.const(0, "float32"), dtype="float32"
+ )
+ gv: R.Tuple(R.Tensor((1, 128, 128), dtype="float32")) = (lv,)
+ R.output(gv)
+ return gv
+
+ example_args = (torch.randn(1, 128, 128, dtype=torch.float32),)
+ verify_model(NewZeros(), example_args, {}, expected1)
+
+
def test_to_copy():
# float
class ToFloat(Module):
diff --git a/tests/python/relax/test_frontend_from_fx.py b/tests/python/relax/test_frontend_from_fx.py
index fdec5ed19c..0d50b4a112 100644
--- a/tests/python/relax/test_frontend_from_fx.py
+++ b/tests/python/relax/test_frontend_from_fx.py
@@ -3352,6 +3352,31 @@ def test_new_ones():
verify_model(NewOnes(), input_info, {}, expected1)
+def test_new_zeros():
+ input_info = [([1, 128, 128], "float32")]
+
+ class NewZeros(Module):
+ def forward(self, x):
+ return x.new_zeros(1, 128, 128)
+
+ @tvm.script.ir_module
+ class expected:
+ @R.function
+ def main(
+ x: R.Tensor((1, 128, 128), dtype="float32")
+ ) -> R.Tensor((1, 128, 128), dtype="float32"):
+ # block 0
+ with R.dataflow():
+ lv: R.Tensor((1, 128, 128), dtype="float32") = R.full(
+ (1, 128, 128), R.const(0.0, "float32"), dtype="float32"
+ )
+ gv: R.Tensor((1, 128, 128), dtype="float32") = lv
+ R.output(gv)
+ return gv
+
+ verify_model(NewZeros(), input_info, {}, expected)
+
+
def test_expand():
input_info = [([1, 2, 3, 4], "float32")]