Set kernel_shape attribute on Conv/ConvTranspose in torchlib

onnxscript/function_libs/torch_lib/ops/core.py

@@ -2305,6 +2305,29 @@ def aten_convolution(
     return result
 
 
+def _conv_kernel_shape(weight: TFloat, complex: bool = False) -> Optional[Sequence[int]]:
+    """Return the spatial (kernel) shape of the convolution weight if statically known.
+
+    The kernel_shape attribute of ONNX Conv/ConvTranspose corresponds to the spatial
+    dimensions of the weight tensor, i.e. all dimensions except the leading two
+    (output and input channels). When ``complex`` is True, the weight has a trailing
+    dimension of size 2 (real/imaginary parts) which is also excluded. Returns None
+    when the spatial shape is not fully static, in which case the attribute is omitted
+    and inferred by ONNX.
+    """
+    shape = weight.shape
+    if shape is None:
+        return None
+    # Exclude the leading output/input channel dims, and for complex weights the
+    # trailing real/imag dim as well.
+    kernel_shape = shape[2:-1] if complex else shape[2:]
+    if len(kernel_shape) == 0:
+        return None
+    if any(not isinstance(dim, int) for dim in kernel_shape):
+        return None
+    return list(kernel_shape)
+
+
 def _aten_convolution_onnx(
     input: TFloat,
     weight: TFloat,
@@ -2328,6 +2351,11 @@ def _aten_convolution_onnx(
     if no_batch:
         input = op.Unsqueeze(input, op.Constant(value_ints=[0]))
 
+    # kernel_shape is the spatial shape of the weight tensor (excluding the
+    # output-channel and input-channel dimensions). It is optional in the ONNX
+    # spec but is set explicitly here to match the ONNX Conv specification.
+    kernel_shape = _conv_kernel_shape(weight)
+
     if transposed:
         result = op.ConvTranspose(
             input,
@@ -2338,6 +2366,7 @@ def _aten_convolution_onnx(
             group=groups,
             dilations=dilations,
             output_padding=output_padding,
+            kernel_shape=kernel_shape,
         )
     else:
         result = op.Conv(
@@ -2348,6 +2377,7 @@ def _aten_convolution_onnx(
             pads=pads,
             group=groups,
             dilations=dilations,
+            kernel_shape=kernel_shape,
         )
 
     if no_batch:
@@ -2387,6 +2417,11 @@ def _aten_convolution_complex_onnx(
     bias_imag = op.Gather(bias, 1, axis=-1)
     bias_zero = op.Expand(op.CastLike(0.0, weight), op.Shape(bias_real))
 
+    # The complex weight has a trailing dimension of size 2 (real/imaginary parts),
+    # so the spatial (kernel) shape excludes both leading channel dimensions and the
+    # trailing complex dimension.
+    kernel_shape = _conv_kernel_shape(weight, complex=True)
+
     if transposed:
         result_real = op.Sub(
             op.ConvTranspose(
@@ -2398,6 +2433,7 @@ def _aten_convolution_complex_onnx(
                 group=groups,
                 dilations=dilations,
                 output_padding=output_padding,
+                kernel_shape=kernel_shape,
             ),
             op.ConvTranspose(
                 input_imag,
@@ -2408,6 +2444,7 @@ def _aten_convolution_complex_onnx(
                 group=groups,
                 dilations=dilations,
                 output_padding=output_padding,
+                kernel_shape=kernel_shape,
             ),
         )
         result_imag = op.Add(
@@ -2420,6 +2457,7 @@ def _aten_convolution_complex_onnx(
                 group=groups,
                 dilations=dilations,
                 output_padding=output_padding,
+                kernel_shape=kernel_shape,
             ),
             op.ConvTranspose(
                 input_imag,
@@ -2430,6 +2468,7 @@ def _aten_convolution_complex_onnx(
                 group=groups,
                 dilations=dilations,
                 output_padding=output_padding,
+                kernel_shape=kernel_shape,
             ),
         )
     else:
@@ -2442,6 +2481,7 @@ def _aten_convolution_complex_onnx(
                 pads=pads,
                 group=groups,
                 dilations=dilations,
+                kernel_shape=kernel_shape,
             ),
             op.Conv(
                 input_imag,
@@ -2451,6 +2491,7 @@ def _aten_convolution_complex_onnx(
                 pads=pads,
                 group=groups,
                 dilations=dilations,
+                kernel_shape=kernel_shape,
             ),
         )
 
@@ -2463,6 +2504,7 @@ def _aten_convolution_complex_onnx(
                 pads=pads,
                 group=groups,
                 dilations=dilations,
+                kernel_shape=kernel_shape,
             ),
             op.Conv(
                 input_imag,
@@ -2472,6 +2514,7 @@ def _aten_convolution_complex_onnx(
                 pads=pads,
                 group=groups,
                 dilations=dilations,
+                kernel_shape=kernel_shape,
             ),
         )
 

tests/function_libs/torch_lib/e2e_ops_tests.py

@@ -679,23 +679,23 @@ def forward(self, x):
         )
         _testing.assert_onnx_program(onnx_program)
 
-        model = onnx_program.model_proto
+        model = onnx_program.model
 
         def _rank(name: str) -> int:
             for vi in (
                 list(model.graph.value_info)
-                + list(model.graph.input)
-                + list(model.graph.output)
+                + list(model.graph.inputs)
+                + list(model.graph.outputs)
             ):
                 if vi.name == name:
                     return len(vi.type.tensor_type.shape.dim)
             raise AssertionError(f"value_info for {name} not found")
 
-        stft_nodes = [n for n in model.graph.node if n.op_type == "STFT"]
+        stft_nodes = [n for n in model.graph if n.op_type == "STFT"]
         self.assertEqual(len(stft_nodes), 1)
         node = stft_nodes[0]
-        signal, frame_step = node.input[0], node.input[1]
-        frame_length = node.input[3]
+        signal, frame_step = node.inputs[0], node.inputs[1]
+        frame_length = node.inputs[3]
         # signal must be rank 3: [batch, signal_length, 1]
         self.assertEqual(_rank(signal), 3)
         # frame_step and frame_length must share the same (scalar) rank
@@ -1266,6 +1266,109 @@ def forward(self, x):
                 got = onnx_program.call_reference({"x": inputs[0]})
                 torch.testing.assert_close(expected, got[0])
 
+    def test_aten_convolution_sets_kernel_shape_when_static(self):
+        # Regression test for https://github.com/microsoft/onnxscript/pull/2972
+        # The emitted Conv node must carry a kernel_shape attribute when the
+        # weight's spatial dimensions are statically known.
+        class Model(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.conv = torch.nn.Conv2d(3, 4, kernel_size=3)
+
+            def forward(self, x):
+                return self.conv(x)
+
+        onnx_program = torch.onnx.export(
+            Model(),
+            (torch.randn(1, 3, 8, 8),),
+            dynamo=True,
+            verbose=False,
+        )
+        _testing.assert_onnx_program(onnx_program)
+
+        conv_nodes = [n for n in onnx_program.model.graph if n.op_type == "Conv"]
+        self.assertEqual(len(conv_nodes), 1)
+        kernel_shape_attrs = [conv_nodes[0].attributes["kernel_shape"]]
+        self.assertEqual(len(kernel_shape_attrs), 1)
+        self.assertEqual(list(kernel_shape_attrs[0].value), [3, 3])
+
+    def test_aten_convolution_transpose_sets_kernel_shape_when_static(self):
+        # Regression test for https://github.com/microsoft/onnxscript/pull/2972
+        # The emitted ConvTranspose node must carry a kernel_shape attribute when
+        # the weight's spatial dimensions are statically known.
+        class Model(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.conv = torch.nn.ConvTranspose2d(3, 4, kernel_size=3)
+
+            def forward(self, x):
+                return self.conv(x)
+
+        onnx_program = torch.onnx.export(
+            Model(),
+            (torch.randn(1, 3, 8, 8),),
+            dynamo=True,
+            verbose=False,
+        )
+        _testing.assert_onnx_program(onnx_program)
+
+        conv_nodes = [n for n in onnx_program.model.graph if n.op_type == "ConvTranspose"]
+        self.assertEqual(len(conv_nodes), 1)
+        self.assertIn("kernel_shape", conv_nodes[0].attributes)
+        kernel_shape_attr = conv_nodes[0].attributes["kernel_shape"]
+        self.assertEqual(list(kernel_shape_attr.value), [3, 3])
+
+    def test_aten_convolution_complex_sets_kernel_shape_when_static(self):
+        # Regression test for https://github.com/microsoft/onnxscript/pull/2972
+        # The complex-conv path must also carry a kernel_shape attribute derived
+        # from the weight's spatial dimensions (excluding the trailing real/imag
+        # dimension) when they are statically known.
+        class Model(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.weight = torch.nn.Parameter(
+                    torch.randn(4, 3, 3, 3, dtype=torch.complex64)
+                )
+
+            def forward(self, x):
+                return torch.nn.functional.conv2d(x, self.weight)
+
+        onnx_program = torch.onnx.export(
+            Model(),
+            (torch.randn(1, 3, 8, 8, dtype=torch.complex64),),
+            dynamo=True,
+            verbose=False,
+        )
+        _testing.assert_onnx_program(onnx_program)
+
+        conv_nodes = [n for n in onnx_program.model.graph if n.op_type == "Conv"]
+        self.assertGreater(len(conv_nodes), 0)
+        for node in conv_nodes:
+            self.assertIn("kernel_shape", node.attributes)
+            kernel_shape_attr = node.attributes["kernel_shape"]
+            self.assertEqual(list(kernel_shape_attr.value), [3, 3])
+
+    def test_aten_convolution_omits_kernel_shape_when_dynamic(self):
+        # Regression test for https://github.com/microsoft/onnxscript/pull/2972
+        # The kernel_shape attribute must be omitted when the weight's spatial
+        # dimensions are dynamic, so that ONNX infers it from the weight tensor.
+        class Model(torch.nn.Module):
+            def forward(self, x, weight):
+                return torch.nn.functional.conv2d(x, weight)
+
+        onnx_program = torch.onnx.export(
+            Model(),
+            (torch.randn(1, 3, 8, 8), torch.randn(4, 3, 3, 3)),
+            dynamo=True,
+            verbose=False,
+            dynamic_shapes=({}, {2: "kh", 3: "kw"}),
+        )
+        _testing.assert_onnx_program(onnx_program)
+
+        conv_nodes = [n for n in onnx_program.model.graph if n.op_type == "Conv"]
+        self.assertEqual(len(conv_nodes), 1)
+        self.assertNotIn("kernel_shape", conv_nodes[0].attributes)
+
     @unittest.skip("see https://github.com/pytorch/pytorch/issues/174668")
     def test_aten_histc_float16(self):
         class Model(torch.nn.Module):