[fx] add acc_ops.expand(_as) and shape inference (#31)

Summary:
Pull Request resolved: https://github.com/pytorch/fx2trt/pull/31

1. acc_ops
    a. 'call_method' `torch.Tensor.expand` will map to `acc_op.expand`
    b. 'call_method' `torch.Tensor.expand_as` will map to `acc_op.expand`
2. shape inference will run for `acc_op.expand`
3. acc_normalizer will map `acc_op.expand` to `acc_op.tile` using the required shape information

Reviewed By: 842974287

Differential Revision: D34565863

fbshipit-source-id: ad94694e548f5b03af9cf3ef8e7c35e236bb0745

Author: alexbeloi

tracer/acc_tracer/acc_ops.py

@@ -2192,3 +2192,27 @@ def gather(*, input, dim, index, sparse_grad=False):
 @register_acc_op
 def index_select(*, input, dim, index):
     return torch.index_select(input, dim, index)
+
+
+@register_custom_acc_mapper_fn(
+    op_and_target=("call_method", "expand_as"),
+    arg_replacement_tuples=[
+        ("input", "input"),
+        ("other", "other"),
+    ]
+)
+def expand_as_mapper(node: torch.fx.Node, _: nn.Module) -> torch.fx.Node:
+    """
+    Maps expand_as(other) to expand(other.size())
+    """
+    with node.graph.inserting_before(node):
+        size_node = node.graph.call_function(
+            size, kwargs={"input": node.kwargs["other"]}
+        )
+        size_node.meta["type"] = torch.Size
+
+        expand_node = node.graph.call_function(
+            expand, kwargs={"input": node.kwargs["input"], "sizes": size_node}
+        )
+        expand_node.meta = node.meta.copy()
+        return expand_node