Add support for aten::dot (apache#9893)

* Add support for aten::dot This implements dot product as a composite of of multiply + sum * address comments Co-authored-by: driazati <[email protected]>
ylc · Feb 16, 2022 · d7e2dbb · d7e2dbb
1 parent ba6cd7a
commit d7e2dbb
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 1 deletion.
diff --git a/python/tvm/relay/frontend/common.py b/python/tvm/relay/frontend/common.py
@@ -958,7 +958,7 @@ def try_resolve_var_to_const(x, graph_params):
 
 
 def set_span(sym, node_name):
-    """Set up the sapn of relay expression(s) while converting OP"""
+    """Set up the span of relay expression(s) while converting OP"""
 
     class SpanFiller(ExprMutator):
         """SpanFiller"""

diff --git a/python/tvm/relay/frontend/pytorch.py b/python/tvm/relay/frontend/pytorch.py
@@ -2852,6 +2852,10 @@ def einsum(self, inputs, input_types):
         equation, data = inputs
         return _op.einsum(data, equation)
 
+    def dot(self, inputs, _):
+        lhs, rhs = inputs
+        return _op.sum(_op.multiply(lhs, rhs))
+
     # Operator mappings
     def create_convert_map(self):
         self.convert_map = {
@@ -3076,6 +3080,7 @@ def create_convert_map(self):
             "aten::bucketize": self.bucketize,
             "aten::roll": self.roll,
             "aten::einsum": self.einsum,
+            "aten::dot": self.dot,
         }
 
     def update_convert_map(self, custom_map):

diff --git a/tests/python/frontend/pytorch/test_forward.py b/tests/python/frontend/pytorch/test_forward.py
@@ -4087,5 +4087,14 @@ def test_fn(equation):
     verify_model(test_fn("ij,jk,km->im"), [x, y, z])
 
 
+@tvm.testing.uses_gpu
+def test_dot():
+    def test_fn(x):
+        return x.dot(x)
+
+    x = torch.randn([4])
+    verify_model(test_fn, [x])
+
+
 if __name__ == "__main__":
     pytest.main([__file__])