[Relax][PyTorch] Add support for eye op in fx graph (apache#17908)

Author: tlopex

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,

tests/python/relax/test_frontend_from_fx.py

@@ -5342,5 +5342,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()