Repository: microsoft/onnxscript

Issue:
ops: convolution
Implement aten::convolution


Discussion:

Resolving patch:
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
@@ -1068,39 +1068,19 @@ def aten_conv2d(
         zero = op.CastLike(0.0, input)
         bias = op.Expand(zero, bias_shape)
 
-    result = _aten_conv2d_onnx(
-        input, weight, bias, strides=strides, pads=pads, dilations=dilations, groups=groups
+    result = _aten_convolution_onnx(
+        input,
+        weight,
+        bias,
+        strides=strides,
+        pads=pads,
+        dilations=dilations,
+        groups=groups,
     )
 
     return result
 
 
-@torch_op("aten::conv2d", overload=True)
-def _aten_conv2d_onnx(
-    input: TFloat,
-    weight: TFloat,
-    bias: TFloat,
-    strides: Sequence[int],
-    pads: Sequence[int],
-    dilations: Sequence[int],
-    groups: int,
-) -> TFloat:
-    """conv2d with attributes pre-computed to fit the ONNX spec."""
-
-    # Torch input may be 3D or 4D, ONNX requires 4D
-    no_batch = op.Size(op.Shape(input)) == 3
-    if no_batch:
-        input = op.Unsqueeze(input, op.Constant(value_ints=[0]))
-
-    result = op.Conv(
-        input, weight, bias, strides=strides, pads=pads, group=groups, dilations=dilations
-    )
-
-    if no_batch:
-        result = op.Squeeze(result, op.Constant(value_ints=[0]))
-    return result
-
-
 def aten_conv3d(
     input: TensorType,
     weight: TensorType,
@@ -1146,20 +1126,99 @@ def aten_conv_transpose1d(
     raise NotImplementedError()
 
 
+@torch_op("aten::convolution", trace_only=True)
 def aten_convolution(
-    input: TensorType,
-    weight: TensorType,
-    bias: Optional[TensorType],
+    input: TFloat,
+    weight: TFloat,
+    bias: Optional[TFloat],
     stride: Sequence[int],
-    padding: INT64,
+    padding: Sequence[int],
     dilation: Sequence[int],
     transposed: bool,
-    output_padding: INT64,
+    output_padding: Sequence[int],
     groups: int,
-) -> TensorType:
+) -> TFloat:
     # convolution(Tensor input, Tensor weight, Tensor? bias, int[] stride, SymInt[] padding, int[] dilation, bool transposed, SymInt[] output_padding, int groups) -> Tensor
 
-    raise NotImplementedError()
+    if not isinstance(padding, Sequence):
+        padding = (padding, padding)
+    pads = [*padding, *padding]
+
+    if not isinstance(dilation, Sequence):
+        dilation = (dilation, dilation)
+    dilations = list(dilation)
+
+    if not isinstance(stride, Sequence):
+        stride = (stride, stride)
+    strides = list(stride)
+
+    if bias is None:
+        weight_dim_0 = op.Shape(weight, start=0, end=1)
+        bias_shape = op.Expand(weight_dim_0, op.Constant(value_ints=[1]))
+        zero = op.CastLike(0.0, input)
+        bias = op.Expand(zero, bias_shape)
+
+    result = _aten_convolution_onnx(
+        input,
+        weight,
+        bias,
+        strides=strides,
+        pads=pads,
+        dilations=dilations,
+        transposed=transposed,
+        output_padding=output_padding,
+        groups=groups,
+    )
+
+    return result
+
+
+@torch_op("aten::convolution", overload=True)
+def _aten_convolution_onnx(
+    input: TFloat,
+    weight: TFloat,
+    bias: TFloat,
+    strides: Sequence[int],
+    pads: Sequence[int],
+    dilations: Sequence[int],
+    transposed: bool = False,
+    output_padding: Sequence[int] = (0,),
+    groups: int = 1,
+) -> TFloat:
+    """ConvXd with attributes pre-computed to fit the ONNX spec."""
+
+    weight_size = op.Size(op.Shape(weight))
+    no_batch = op.Size(op.Shape(input)) != weight_size
+
+    if no_batch:
+        input = op.Unsqueeze(input, op.Constant(value_ints=[0]))
+
+    if transposed:
+        result = op.ConvTranspose(
+            input,
+            weight,
+            bias,
+            strides=strides,
+            pads=pads,
+            group=groups,
+            dilations=dilations,
+            output_padding=output_padding,
+        )
+    else:
+        result = op.Conv(
+            input,
+            weight,
+            bias,
+            strides=strides,
+            pads=pads,
+            group=groups,
+            dilations=dilations,
+        )
+
+    if no_batch:
+        result = op.Squeeze(result, op.Constant(value_ints=[0]))
+
+    return result
 
 
 def aten_convolution_backward(


Test patch:
diff --git a/onnxscript/tests/function_libs/torch_aten/extra_opinfo.py b/onnxscript/tests/function_libs/torch_aten/extra_opinfo.py
new file mode 100644
--- /dev/null
+++ b/onnxscript/tests/function_libs/torch_aten/extra_opinfo.py
@@ -0,0 +1,114 @@
+"""
+Test data for aten operators which don't exist in PyTorch file:
+pytorch/torch/testing/_internal/common_methods_invocations.py.
+"""
+
+import functools
+from typing import Any, List
+
+import torch
+from torch import testing as torch_testing
+from torch.testing._internal import common_dtype, common_utils
+from torch.testing._internal.opinfo import core as opinfo_core
+
+
+def sample_inputs_convolution(op_info, device, dtype, requires_grad, **kwargs):
+    del op_info
+    make_arg = functools.partial(
+        torch_testing.make_tensor, device=device, dtype=dtype, requires_grad=requires_grad
+    )
+
+    # Ordered as shapes for input, weight, bias,
+    # and a dict of values of (stride, padding, dilation, groups)
+    cases: tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...], dict[str, Any]] = (  # type: ignore[assignment]
+        (
+            (1, 3, 4),
+            (3, 3, 3),
+            (3,),
+            {
+                "stride": (2,),
+                "padding": (2,),
+                "dilation": (1,),
+                "transposed": False,
+                "output_padding": (0,),
+                "groups": 1,
+            },
+        ),
+        (
+            (1, 3, 4),
+            (3, 3, 3),
+            None,
+            {
+                "stride": (2,),
+                "padding": (2,),
+                "dilation": (1,),
+                "transposed": True,
+                "output_padding": (0,),
+                "groups": 1,
+            },
+        ),
+        (
+            (1, 3, 224, 224),
+            (32, 3, 3, 3),
+            None,
+            {
+                "stride": (2, 2),
+                "padding": (1, 1),
+                "dilation": (1, 1),
+                "transposed": False,
+                "output_padding": (0, 0),
+                "groups": 1,
+            },
+        ),
+        # FIXME(jiz): Uncomment out these test data once
+        # torch 2.0 is released.
+        # (
+        #     (1, 3, 224, 224, 224),
+        #     (32, 3, 3, 3, 3),
+        #     (32,),
+        #     {
+        #         "stride": (2, 2, 2),
+        #         "padding": (1, 1, 1),
+        #         "dilation": (1, 1, 1),
+        #         "transposed": False,
+        #         "output_padding": (0, 0, 0),
+        #         "groups": 1,
+        #     },
+        # ),
+        (
+            (2, 4, 6, 6),
+            (4, 1, 3, 3),
+            (4,),
+            {
+                "stride": (3, 2),
+                "padding": (1, 1),
+                "dilation": (1, 1),
+                "transposed": True,
+                "output_padding": (0, 0),
+                "groups": 4,
+            },
+        ),
+    )
+
+    for input_shape, weight, bias, kwargs in cases:  # type: ignore[assignment]
+        yield opinfo_core.SampleInput(
+            make_arg(input_shape),
+            args=(make_arg(weight), make_arg(bias) if bias is not None else bias),
+            kwargs=kwargs,
+        )
+
+
+OP_DB: List[opinfo_core.OpInfo] = [
+    opinfo_core.OpInfo(
+        "convolution",
+        aliases=("convolution",),
+        aten_name="convolution",
+        dtypes=common_dtype.floating_and_complex_types_and(torch.int64, torch.bfloat16),
+        sample_inputs_func=sample_inputs_convolution,
+        supports_forward_ad=True,
+        supports_fwgrad_bwgrad=True,
+        gradcheck_nondet_tol=common_utils.GRADCHECK_NONDET_TOL,
+        skips=(),
+        supports_out=False,
+    ),
+]
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
@@ -20,6 +20,7 @@
 from onnxscript.function_libs.torch_aten.ops import nn as nn_ops
 from onnxscript.function_libs.torch_aten.ops import special as special_ops
 from onnxscript.tests.common import version_utils
+from onnxscript.tests.function_libs.torch_aten import extra_opinfo
 
 T = TypeVar("T")
 
@@ -168,6 +169,9 @@ def duplicate_opinfo(opinfos: list[opinfo_core.OpInfo], name: str, new_names: tu
 # Create a copy of the op_db to modify
 OPS_DB = copy.deepcopy(common_methods_invocations.op_db)
 
+# Append extra op_db into the op database for testing
+OPS_DB.extend(extra_opinfo.OP_DB)
+
 # Modify this section ##########################################################
 
 
@@ -357,10 +361,11 @@ def _where_input_wrangler(
     "argmax": core_ops.aten_argmax,
     "argmin": core_ops.aten_argmin,
     "clamp": core_ops.aten_clamp,
-    "nn.functional.gelu": nn_ops.aten_gelu,
+    "convolution": core_ops.aten_convolution,
     "index_select": core_ops.aten_index_select,
     "native_layer_norm": core_ops.aten_native_layer_norm,
     "nn.functional.conv2d": core_ops.aten_conv2d,
+    "nn.functional.gelu": nn_ops.aten_gelu,
     "nn.functional.linear": nn_ops.aten_linear,
     "slice": core_ops.aten_slice,
     "sum": (core_ops.aten_sum_dim_IntList, _sum_input_wrangler),