Repository: microsoft/onnxscript

Issue:
Compute the correct number of output for Split | fix(graph_building)
When number of outputs is known, Split should not a sequence of tensors.


Discussion:
## [Codecov](https://codecov.io/gh/microsoft/onnx-script/pull/622?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=microsoft) Report
> Merging [#622](https://codecov.io/gh/microsoft/onnx-script/pull/622?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=microsoft) (1889ec7) into [main](https://codecov.io/gh/microsoft/onnx-script/commit/f1591c7e1879ce5f595be5a99b868820fe64fb1a?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=microsoft) (f1591c7) will **decrease** coverage by `0.02%`.
> The diff coverage is `33.33%`.

```diff
@@            Coverage Diff             @@
##             main     #622      +/-   ##
==========================================
- Coverage   74.20%   74.19%   -0.02%     
==========================================
  Files         107      107              
  Lines       11302    11305       +3     
  Branches     1177     1178       +1     
==========================================
+ Hits         8387     8388       +1     
- Misses       2592     2593       +1     
- Partials      323      324       +1     
```


| [Impacted Files](https://codecov.io/gh/microsoft/onnx-script/pull/622?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=microsoft) | Coverage Δ | |
|---|---|---|
| [...xscript/function\_libs/torch\_aten/graph\_building.py](https://codecov.io/gh/microsoft/onnx-script/pull/622?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=microsoft#diff-b25ueHNjcmlwdC9mdW5jdGlvbl9saWJzL3RvcmNoX2F0ZW4vZ3JhcGhfYnVpbGRpbmcucHk=) | `77.43% <33.33%> (-0.53%)` | :arrow_down: |

Help us with your feedback. Take ten seconds to tell us [how you rate us](https://about.codecov.io/nps?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=microsoft). Have a feature suggestion? [Share it here.](https://app.codecov.io/gh/feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=microsoft)

Looks reasonable to me. But will leave it to those more familiar with the exporter comment/review.

Resolving patch:
diff --git a/onnxscript/evaluator.py b/onnxscript/evaluator.py
--- a/onnxscript/evaluator.py
+++ b/onnxscript/evaluator.py
@@ -316,10 +316,12 @@ def _rename_io(prefix, i, arg):
     return f"{prefix}{i}"
 
 
-def _compute_num_outputs(schema: onnx.defs.OpSchema, *args: Any, **kwargs: Any):
-    """Returns the number of outputs expected.
-    TODO: Use ONNX type inference to replace the special-case handling below.
-    """
+def compute_num_outputs(
+    schema: onnx.defs.OpSchema, args: Sequence[Any], kwargs: Mapping[str, Any]
+) -> int:
+    """Returns the number of outputs expected."""
+
+    # TODO: Use ONNX type inference to replace the special-case handling below.
     if schema.domain == "":
         if schema.name == "BatchNormalization":
             if not kwargs.get("training_mode", 0):
@@ -414,7 +416,7 @@ def _call_ort(
     # Construct ONNX model with a single op call:
     inputs = [_rename_io("input", i, arg) for i, arg in enumerate(args)]
 
-    num_outputs = _compute_num_outputs(schema, *args, **kwargs)
+    num_outputs = compute_num_outputs(schema, args, kwargs)
     outputs = [f"output{i}" for i in range(num_outputs)]
 
     node = onnx.helper.make_node(schema.name, inputs, outputs, domain=schema.domain, **kwargs)
diff --git a/onnxscript/function_libs/torch_lib/graph_building.py b/onnxscript/function_libs/torch_lib/graph_building.py
--- a/onnxscript/function_libs/torch_lib/graph_building.py
+++ b/onnxscript/function_libs/torch_lib/graph_building.py
@@ -458,11 +458,12 @@ def add_op_call(
         onnx_attributes: Mapping[str, ValidArgumentType],
     ):
         # Compute outputs from the onnx_op op schema
+        n_outputs = evaluator.compute_num_outputs(onnx_op_schema, onnx_inputs, onnx_attributes)
         result = self._add_torchscript_op_call(
             f"onnx::{onnx_op_schema.name}",
             onnx_inputs,
             onnx_attributes,
-            n_outputs=len(onnx_op_schema.outputs),
+            n_outputs=n_outputs,
         )
 
         return result
diff --git a/onnxscript/function_libs/torch_lib/ops/core.py b/onnxscript/function_libs/torch_lib/ops/core.py
--- a/onnxscript/function_libs/torch_lib/ops/core.py
+++ b/onnxscript/function_libs/torch_lib/ops/core.py
@@ -1039,9 +1039,9 @@ def aten_choose_qparams_optimized(
 
 
 @torch_op("aten::chunk")
-def aten_chunk(self: TTensor, chunks: INT64, dim: int = 0) -> TTensor:
+def aten_chunk(self: TTensor, chunks: int, dim: int = 0) -> Sequence[TTensor]:
     """chunk(Tensor(a -> *) self, int chunks, int dim=0) -> Tensor(a)[]"""
-
+    # This will create a Sequence of tensors
     neg_1 = op.Constant(value_ints=[-1])
     # Get size of specified dim
     self_shape = op.Shape(self)
@@ -1059,7 +1059,8 @@ def aten_chunk(self: TTensor, chunks: INT64, dim: int = 0) -> TTensor:
     if remainder > 0:  # type: ignore[operator]
         # Append the remainder to the [n, n, n, n, ..., r]
         list_split = op.Concat(list_split, op.Reshape(remainder, neg_1), axis=0)
-    return op.Split(self, list_split, axis=dim)
+
+    return op.SplitToSequence(self, list_split, axis=dim)
 
 
 @torch_op("aten::clamp", trace_only=True)


Test patch:
diff --git a/onnxscript/tests/function_libs/torch_lib/ops_test.py b/onnxscript/tests/function_libs/torch_lib/ops_test.py
--- a/onnxscript/tests/function_libs/torch_lib/ops_test.py
+++ b/onnxscript/tests/function_libs/torch_lib/ops_test.py
@@ -193,8 +193,8 @@ def run_test_output_match(
                 torch_output = op(*inputs, **cpu_sample.kwargs)
 
                 reference_torch_outputs, _ = pytree.tree_flatten(torch_output)
-                if op.name.startswith("split"):
-                    # Hack for handling split
+                if op.name.startswith("split") or op.name.startswith("chunk"):
+                    # Hack for handling split and chunk
                     # Split returns a Sequence that should be treats as a single
                     # value. So we wrap it into a tuple.
                     # TODO(justinchuby): Find a more general solution
diff --git a/onnxscript/tests/function_libs/torch_lib/ops_test_data.py b/onnxscript/tests/function_libs/torch_lib/ops_test_data.py
--- a/onnxscript/tests/function_libs/torch_lib/ops_test_data.py
+++ b/onnxscript/tests/function_libs/torch_lib/ops_test_data.py
@@ -602,9 +602,6 @@ def _where_input_wrangler(
         variant_name="partial_views",
         reason="ONNX doesn't have partial view for tensor",
     ),
-    xfail(
-        "chunk", reason="fixme: ORT error", test_class_name="TestOutputConsistencyFullGraph"
-    ),
     xfail("logcumsumexp", reason="naive implementation not numerically stable"),
     xfail(
         "max",