cleanup

Author: mshr-h

python/tvm/relax/frontend/torch/fx_translator.py

@@ -119,23 +119,6 @@ def _retrieve_args(self, node):
         else:
             return node
 
-    @staticmethod
-    def _promote_binary_op_args(lhs, rhs):
-        if isinstance(lhs, relax.Expr) and isinstance(rhs, relax.Expr):
-            return lhs, rhs
-        elif isinstance(lhs, relax.Expr):
-            assert isinstance(lhs.struct_info, relax.TensorStructInfo)
-            return lhs, relax.const(rhs, lhs.struct_info.dtype)
-        elif isinstance(rhs, relax.Expr):
-            assert isinstance(rhs.struct_info, relax.TensorStructInfo)
-            return relax.const(lhs, rhs.struct_info.dtype), rhs
-        else:
-            assert False
-
-    def _call_binary_op(self, op, lhs, rhs):
-        lhs, rhs = TorchFXImporter._promote_binary_op_args(lhs, rhs)
-        return self.block_builder.emit(op(lhs, rhs))
-
     ########## Unary Ops ##########
 
     def _unary_op(self, op: Callable) -> Callable:
@@ -246,17 +229,29 @@ def _binary_op(self, relax_op: Callable, intrinsic_op: Callable) -> Callable:
         from torch import fx
 
         def convert(node: fx.Node) -> relax.Var:
+            def promote_binary_op_args(lhs, rhs):
+                if isinstance(lhs, relax.Expr) and isinstance(rhs, relax.Expr):
+                    return lhs, rhs
+                elif isinstance(lhs, relax.Expr):
+                    assert isinstance(lhs.struct_info, relax.TensorStructInfo)
+                    return lhs, relax.const(rhs, lhs.struct_info.dtype)
+                elif isinstance(rhs, relax.Expr):
+                    assert isinstance(rhs.struct_info, relax.TensorStructInfo)
+                    return relax.const(lhs, rhs.struct_info.dtype), rhs
+                else:
+                    assert False
+
+            def call_binary_op(op, lhs, rhs):
+                lhs, rhs = promote_binary_op_args(lhs, rhs)
+                return self.block_builder.emit(op(lhs, rhs))
+
             lhs, rhs = self.retrieve_args(node)
             if isinstance(lhs, relax.Var) or isinstance(rhs, relax.Var):
-                return self._call_binary_op(relax_op, lhs, rhs)
+                return call_binary_op(relax_op, lhs, rhs)
             elif isinstance(lhs, relax.expr.Constant):
-                return self._call_binary_op(
-                    relax_op, lhs, relax.const(rhs, dtype=lhs.struct_info.dtype)
-                )
+                return call_binary_op(relax_op, lhs, relax.const(rhs, dtype=lhs.struct_info.dtype))
             elif isinstance(rhs, relax.expr.Constant):
-                return self._call_binary_op(
-                    relax_op, relax.const(lhs, dtype=rhs.struct_info.dtype), rhs
-                )
+                return call_binary_op(relax_op, relax.const(lhs, dtype=rhs.struct_info.dtype), rhs)
             return intrinsic_op(lhs, rhs)
 
         return convert