add einsum in pytorch frontend (apache#9651)

* add einsum in pytorch frontend

* add einsum in pytorch frontend

Author: Meteorix

python/tvm/relay/frontend/pytorch.py

@@ -2819,6 +2819,10 @@ def slide_axes(inp, shape, ax):
 
         return out
 
+    def einsum(self, inputs, input_types):
+        equation, data = inputs
+        return _op.einsum(data, equation)
+
     # Operator mappings
     def create_convert_map(self):
         self.convert_map = {
@@ -3063,6 +3067,7 @@ def create_convert_map(self):
             "aten::searchsorted": self.searchsorted,
             "aten::bucketize": self.bucketize,
             "aten::roll": self.roll,
+            "aten::einsum": self.einsum,
         }
 
     def update_convert_map(self, custom_map):

tests/python/frontend/pytorch/test_forward.py

@@ -4020,5 +4020,17 @@ def test_fn(shifts, dims):
     verify_model(test_fn(shifts=(2, 1), dims=(0, 1)), [x])
 
 
+@tvm.testing.uses_gpu
+def test_einsum():
+    def test_fn(equation):
+        return lambda *x: torch.einsum(equation, *x)
+
+    x = torch.ones([2, 3])
+    y = torch.ones([3, 4])
+    z = torch.ones([4, 5])
+    verify_model(test_fn("ij,jk"), [x, y])
+    verify_model(test_fn("ij,jk,km->im"), [x, y, z])
+
+
 if __name__ == "__main__":
     pytest.main([__file__])