From e6c37265109f6dee4c9e61e7b399ba7dc28d1a36 Mon Sep 17 00:00:00 2001 From: Shushi Hong <820958424@qq.com> Date: Mon, 28 Apr 2025 23:18:59 +0800 Subject: [PATCH 1/2] Update fx_translator.py --- python/tvm/relax/frontend/torch/fx_translator.py | 1 + 1 file changed, 1 insertion(+) diff --git a/python/tvm/relax/frontend/torch/fx_translator.py b/python/tvm/relax/frontend/torch/fx_translator.py index 07b3df20aa72..44d5ba4c43a6 100644 --- a/python/tvm/relax/frontend/torch/fx_translator.py +++ b/python/tvm/relax/frontend/torch/fx_translator.py @@ -821,6 +821,7 @@ def create_convert_map( "clone": lambda node: self.env[node.args[0]], "empty": self._empty, "empty_like": self._empty_like, + "eye": self._eye, "fill": self._fill, "fill_": self._inplace_fill, "full": self._full, From 7edd6a8d4b9d29bacfc9b407ff4f6a2316f73762 Mon Sep 17 00:00:00 2001 From: Shushi Hong <820958424@qq.com> Date: Mon, 28 Apr 2025 23:21:55 +0800 Subject: [PATCH 2/2] Update test_frontend_from_fx.py --- tests/python/relax/test_frontend_from_fx.py | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/tests/python/relax/test_frontend_from_fx.py b/tests/python/relax/test_frontend_from_fx.py index f60f158cbfa4..ef8d42abebcb 100644 --- a/tests/python/relax/test_frontend_from_fx.py +++ b/tests/python/relax/test_frontend_from_fx.py @@ -5292,5 +5292,23 @@ def main( verify_model(BFloat16Model(), [([10, 10], "bfloat16"), ([10, 10], "bfloat16")], {}, Expected) +def test_eye(): + import numpy as np + + class Eye(Module): + def forward(self, input): + return torch.eye(3) + + graph_model = fx.symbolic_trace(Eye()) + mod = from_fx(graph_model, [([3, 3], "float32")]) + assert len(mod["main"].body.blocks) == 1 + assert len(mod["main"].body.blocks[0].bindings) == 1 + assert isinstance(mod["main"].body.blocks[0].bindings[0].value, relax.Constant) + tvm.testing.assert_allclose( + mod["main"].body.blocks[0].bindings[0].value.data.numpy(), + np.eye(3, dtype="float32"), + ) + + if __name__ == "__main__": tvm.testing.main()