[ONNX] Support aten::broadcast_to (pytorch#101833)

Support aten::broadcast as the way we support on aten::expand.

Fix pytorch#92678 pytorch#101768
Pull Request resolved: pytorch#101833
Approved by: https://github.com/thiagocrepaldi

Author: titaiwangms

test/onnx/test_op_consistency.py

@@ -56,7 +56,9 @@
     [
         "atan",
         "atan2",
+        "broadcast_to",
         "ceil",
+        "expand",
         "flatten",
         "logical_not",
         "nn.functional.scaled_dot_product_attention",

torch/onnx/symbolic_opset9.py

@@ -52,6 +52,7 @@
     "bitwise_or",
     "bmm",
     "broadcast_tensors",
+    "broadcast_to",
     "bucketize",
     "cat",
     "cdist",
@@ -965,6 +966,27 @@ def expand(g: jit_utils.GraphContext, self, size, implicit):
     return g.op("Expand", self, size)
 
 
+@_onnx_symbolic("aten::broadcast_to")
+@symbolic_helper.quantized_args(True)
+@_beartype.beartype
+def broadcast_to(g: jit_utils.GraphContext, self, size):
+    size = symbolic_helper._maybe_get_const(size, "is")
+    if not symbolic_helper._is_value(size):
+        size = g.op("Constant", value_t=torch.LongTensor(size))
+    elif symbolic_helper._is_packed_list(size):
+        # Expand with -1 dim value means dim is unchanged.
+        # Since onnx::expand supports two-way broadcasting,
+        # -1 dim value can be exported to onnx as 1
+        size = symbolic_helper._reshape_helper(
+            g, stack(g, size, 0), g.op("Constant", value_t=torch.tensor([-1]))
+        )
+    dtype = _type_utils.JitScalarType.INT64
+    ones = ones_like(g, size, dtype)
+    neg_ones = mul(g, ones, g.op("Constant", value_t=torch.tensor(-1)))
+    size = where(g, g.op("Equal", size, neg_ones), ones, size)
+    return g.op("Expand", self, size)
+
+
 @_onnx_symbolic("aten::expand_as")
 @symbolic_helper.quantized_args(True, True)
 @_beartype.beartype