Repository: microsoft/onnxscript

Issue:
Torchscript Graph: Support dynamic axes
Build support for dynamic axes in inputs


Discussion:
TODO: Try make inputs dynamic as default, and static should be specified if it's necessary.
Tried setting all sizes dynamic as default, but it breaks easily in ORT.

The modified in graph_building.py is easy, we can simply utilize setType API when addin input

```python
torchscript_input_value = torch._C.TensorType.create_from_tensor(input_value)
rank = torchscript_input_value.dim()
torch_value.setType(torchscript_input_value.with_sizes([None] * rank))
```

It works ideally when there is no "restricted-shape nodes" involved. like Add,
```python
    def test_dynamic_axes_on_add(self):
        class DynamicAdd(torch.nn.Module):
            def forward(self, x, y):
                return torch.ops.aten.add(x, y)

        x = torch.randn(2, 3)
        y = torch.randn(2, 3)
        input_x = torch.randn(2, 4)
        input_y = torch.randn(2, 4)
        _run_test_with_fx_to_onnx_exporter_reference_runtime(
            DynamicAdd(), (x, y), additional_test_inputs=[(input_x, input_y)]
        )
```

However, with following model, it breaks in unexpected Expand nodes that decomposed module generated during make_fx, it doesn't work even if I only provide certain dims to be dynamic. (Should be something we have to figure out.)
```python
    def test_dynamic_axes_on_matmul(self):
        class DynamicMatMul(torch.nn.Module):
            def forward(self, x, y):
                return torch.ops.aten.matmul(x, y)

        x = torch.randn(2, 3, 6)
        y = torch.randn(2, 6, 4)
        input_x = torch.randn(2, 3, 4)
        input_y = torch.randn(2, 4, 4)
        # self.run_test(
        #     DynamicMatMul(),
        #     (x, y),
        #     input_names=["x", "y"],
        #     dynamic_axes={"x": {2: "length"}, "y": {1: "length"}},
        #     additional_test_inputs=[(input_x, input_y)],
        # )
        _run_test_with_fx_to_onnx_exporter_reference_runtime(
            DynamicMatMul(), (x, y), additional_test_inputs=[(input_x, input_y)]
        )
```
with error message:
```
onnxruntime.capi.onnxruntime_pybind11_state.InvalidArgument: [ONNXRuntimeError] : 2 : INVALID_ARGUMENT : Non-zero status code returned while running Expand node. Name:'_0x8a230f0_n1' Status Message: invalid expand shape
```

I figure it'd be safe we do it the old way (set static as default), since static shape is still the safest exporting, so I will be working on adding dynamic_axes for API.

cc @BowenBao if you have a good idea on implementing dynamic axes n new exporter.
> However, with following model, it breaks in unexpected Expand nodes that decomposed module generated during make_fx,

Does it pass onnx checker full check and onnx reference evaluator?

For shape/type in general, I'd encourage carrying over what fx graph had, and validating & amending missing ones by onnx shape/type inference.

If dynamo export w/ symbolic tracing mode works well, we should get a fx graph full of shapes and dtypes, with symbolic dimension sizes set on a as needed basis. afaik it is supported via sympy. @jcwchen please offer insight on onnx side how we could utilize that info with onnx shape inference.

If this works well we can leave `dynamic axes` in the past.
I don't have fancy tips for deciding whether to use symbolic dimension or not. I think it's really operator-by-operator. ~Take Expand as an example [here](https://github.com/microsoft/onnxruntime/blob/main/onnxruntime/core/providers/cpu/tensor/expand.cc#L63): it cannot accept a symbolic input shape during Runtime, although onnx shape_inference is fine with it (won't throw an error for this symbolic case).~ Just talked to Ti-Tai offline, it turns out that the Expand input shape needs certain shape due to static Expand's size (Constant). This issue could be resolved by changing Expand's size from the exporter.

Sorry, I don't understand the issue. We can distinguish between symbolic shapes used in static-shape-inference and shapes specified as real-inputs to certain ops like Expand or Reshape. The two are not the same. I guess the issue is in the exporter when running in a mode like trace-mode (or whatever it is called, I am not sure about the name used in the exporter) about how to generalize for inputs of other shapes?
It seems to me the problems would arise when a pytorch model invokes Expand, and we don't have a graph-representation of the computation for the second input of Expand (namely, the target output shape), but we have only a constant tensor. Is that the concern/issue?
Ahhh, I get it now. We messed up the Expand conversion. I see we used to do constant 1 in size to represent unchange dim, maybe that's the key.
@titaiwangms could you post the "matmul w/ expand" onnx graph for illustration?

I agree with Rama's theory. Maybe this particular issue is due to statically constructing the symbols as a constant tensor.


![expand_matmul](https://user-images.githubusercontent.com/18010845/221731904-6f079490-8393-40a5-b52a-5ef0a95ca2f3.png)

Yes, @jcwchen just helped me figure that out. So even though I set all input sizes dynamic in torchscript, the Expand node we have right now does not take dynamic shape into consideration like we used to in old exporter:
https://github.com/microsoft/onnx-script/blob/925e81cc1736c7766e336cc470f23be432929aca/onnxscript/function_libs/torch_aten/ops/core.py#L2037-L2042

What we had
https://github.com/pytorch/pytorch/blob/b818b3fe1c1fa90529b8500cd8ef800bec8415e8/torch/onnx/symbolic_opset9.py#L909-L927

Resolving patch:
diff --git a/onnxscript/function_libs/torch_aten/graph_building.py b/onnxscript/function_libs/torch_aten/graph_building.py
--- a/onnxscript/function_libs/torch_aten/graph_building.py
+++ b/onnxscript/function_libs/torch_aten/graph_building.py
@@ -313,21 +313,29 @@ def initializers(self) -> Mapping[str, torch.Tensor]:
 
     @beartype
     def add_input(
-        self, input_name: str, input_value: Optional[torch.Tensor] = None
+        self,
+        input_name: Optional[str],
+        shape: Optional[Union[torch.Size, Sequence[Union[int, str, None]]]] = None,
+        dtype: Optional[torch.dtype] = None,
     ) -> TorchScriptTensor:
-        # TODO: Take in a TorchScriptTensor?
-        # TODO: Support dynamic shapes
-        if input_value is None:
+        if input_name is None:
             # This input argument is None, which is mapped
             # to a NULL value in TorchScript type system.
             torch_value = _create_op_call_in_torch_graph(
                 self._torch_graph, "prim::Constant", inputs=(), attributes={}
             )[0]
             torch_value.setType(torch.OptionalType.ofTensor())
-            tensor_value = _wrap_torch_value_to_tensor(torch_value)
-            return tensor_value  # type: ignore[return-value]
-        torch_value = self._torch_graph.addInput(input_name)
-        torch_value.setType(torch.TensorType.create_from_tensor(input_value))
+        else:
+            torch_value = self._torch_graph.addInput(input_name)
+            torch_value.setType(torch_value.type().with_dtype(dtype))  # type: ignore[arg-type]
+            # TODO(titaiwang): This approach loses the information that "same SymInts
+            # indicates same shape", for example, [symint0, symint0, symint1]
+            # would all be [None, None, None]
+            torch_value.setType(
+                torch_value.type().with_sizes(
+                    [dim if isinstance(dim, int) else None for dim in shape]  # type: ignore[union-attr]
+                )
+            )
         tensor_value = _wrap_torch_value_to_tensor(torch_value)
         return tensor_value  # type: ignore[return-value]
 
@@ -440,7 +448,6 @@ def add_op_call(
         onnx_attributes: Mapping[str, ValidArgumentType],
     ):
         # Compute outputs from the onnx_op op schema
-
         result = self._add_torchscript_op_call(
             f"onnx::{onnx_op_schema.name}",
             onnx_inputs,
@@ -461,7 +468,6 @@ def add_function_call(
         self._function_store[identifier] = onnx_function
 
         # Compute outputs from the function schema
-
         result = self._add_torchscript_op_call(
             f"{onnx_function.function_ir.domain}::{onnx_function.name}",
             onnx_inputs,
@@ -511,7 +517,6 @@ def to_model_proto(self, opset_version: int) -> onnx.ModelProto:
                 onnxscript.proto2text(onnx_model),
                 self.torch_graph,
             )
-
         return onnx_model
 
     def apply(self, graph_pass: Callable, *args, **kwargs) -> None:
diff --git a/onnxscript/function_libs/torch_aten/ops/core.py b/onnxscript/function_libs/torch_aten/ops/core.py
--- a/onnxscript/function_libs/torch_aten/ops/core.py
+++ b/onnxscript/function_libs/torch_aten/ops/core.py
@@ -5473,6 +5473,19 @@ def aten_swapdims(self: TensorType, dim0: int, dim1: int) -> TensorType:
     raise NotImplementedError()
 
 
+@torch_op("aten::sym_size")
+def aten_sym_size(self: TReal, dim: int = 0) -> TReal:
+    """sym_size(Tensor self, int dim) -> Tensor"""
+    # NOTE: onnx-script doesn't support attribute process,
+    # so op.Shape(self, start=dim, end=dim + 1) is not supported.
+    shape = op.Shape(self)
+    # Reshape helps dim from int to tensor, and
+    # input arguments support attribute processing.
+    start = op.Reshape(dim, [1])
+    end = op.Reshape(dim + 1, [1])
+    return op.Slice(shape, start, end)
+
+
 def aten_symeig(
     self: TensorType, eigenvectors: bool = False, upper: bool = True
 ) -> tuple[TensorType, TensorType]:


Test patch:
diff --git a/onnxscript/function_libs/torch_aten/graph_building_test.py b/onnxscript/function_libs/torch_aten/graph_building_test.py
--- a/onnxscript/function_libs/torch_aten/graph_building_test.py
+++ b/onnxscript/function_libs/torch_aten/graph_building_test.py
@@ -42,7 +42,8 @@ def expected_model():
     def test_traced_graph_on_single_node_is_same_as_compiled_graph(self):
         aten_relu = ops.nn.aten_relu
 
-        x = self.onnxscript_graph.add_input("x", torch.ones((1, 2, 3), dtype=torch.float32))
+        x_tensor = torch.ones((1, 2, 3), dtype=torch.float32)
+        x = self.onnxscript_graph.add_input("x", x_tensor.shape, x_tensor.dtype)
         with evaluator.default_as(self.tracer):
             output = aten_relu(x)
 
@@ -61,7 +62,8 @@ def expected_model(x: FLOAT[1, 2, 3]):
     def test_traced_graph_on_single_node_multi_output_is_same_as_compiled_graph(self):
         aten_topk = ops.core.aten_topk
 
-        x = self.onnxscript_graph.add_input("x", torch.ones((1, 2, 3), dtype=torch.float32))
+        x_tensor = torch.ones((1, 2, 3), dtype=torch.float32)
+        x = self.onnxscript_graph.add_input("x", x_tensor.shape, x_tensor.dtype)
         with evaluator.default_as(self.tracer):
             output = aten_topk(x, 2)
 
diff --git a/onnxscript/tests/function_libs/torch_aten/ops_correctness_test.py b/onnxscript/tests/function_libs/torch_aten/ops_correctness_test.py
--- a/onnxscript/tests/function_libs/torch_aten/ops_correctness_test.py
+++ b/onnxscript/tests/function_libs/torch_aten/ops_correctness_test.py
@@ -971,7 +971,11 @@ def _capture_graph_and_evaluate_torch_script_evaluator(function: Callable, args,
         for i, arg in enumerate(args):
             if isinstance(arg, np.ndarray):
                 input_name = f"input_{i}"
-                input = onnxscript_graph.add_input(input_name, torch.tensor(arg))
+                input = onnxscript_graph.add_input(
+                    input_name,
+                    torch.tensor(arg).shape,
+                    torch.tensor(arg).dtype,
+                )
                 input.value = arg
                 onnxscript_args.append(input)
                 ort_inputs[input_name] = arg
@@ -980,7 +984,11 @@ def _capture_graph_and_evaluate_torch_script_evaluator(function: Callable, args,
                 for j, subarg in enumerate(arg):
                     if isinstance(subarg, np.ndarray):
                         input_name = f"input_{i}_{j}"
-                        input = onnxscript_graph.add_input(input_name, torch.tensor(subarg))
+                        input = onnxscript_graph.add_input(
+                            input_name,
+                            torch.tensor(subarg).shape,
+                            torch.tensor(subarg).dtype,
+                        )
                         input.value = subarg
                         sequence_input.append(input)
                         ort_inputs[input_name] = subarg
@@ -989,7 +997,11 @@ def _capture_graph_and_evaluate_torch_script_evaluator(function: Callable, args,
                 onnxscript_args.append(arg)
         for key, value in kwargs.items():
             if isinstance(value, np.ndarray):
-                input = onnxscript_graph.add_input(key, torch.tensor(value))
+                input = onnxscript_graph.add_input(
+                    key,
+                    torch.tensor(value).shape,
+                    torch.tensor(value).dtype,
+                )
                 input.value = value
                 ort_inputs[key] = value
                 onnxscript_kwargs[key] = input