From d20f0ebe0d365a0bfc7693d3eeb075ada40a42c1 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Thu, 23 Jul 2026 22:02:06 +0000 Subject: [PATCH 1/5] Initial plan From f9a126c11e394c8bc6ff2132b260eb3a8f0750db Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Thu, 23 Jul 2026 22:09:11 +0000 Subject: [PATCH 2/5] Add kernel_shape attribute to op.Conv/ConvTranspose in torchlib --- .../function_libs/torch_lib/ops/core.py | 42 +++++++++++++++++++ 1 file changed, 42 insertions(+) diff --git a/onnxscript/function_libs/torch_lib/ops/core.py b/onnxscript/function_libs/torch_lib/ops/core.py index adf1bad4b6..805ab1ecee 100644 --- a/onnxscript/function_libs/torch_lib/ops/core.py +++ b/onnxscript/function_libs/torch_lib/ops/core.py @@ -2305,6 +2305,23 @@ def aten_convolution( return result +def _conv_kernel_shape(weight: TFloat) -> 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). 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 or len(shape) <= 2: + return None + kernel_shape = shape[2:] + 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 +2345,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 +2360,7 @@ def _aten_convolution_onnx( group=groups, dilations=dilations, output_padding=output_padding, + kernel_shape=kernel_shape, ) else: result = op.Conv( @@ -2348,6 +2371,7 @@ def _aten_convolution_onnx( pads=pads, group=groups, dilations=dilations, + kernel_shape=kernel_shape, ) if no_batch: @@ -2387,6 +2411,16 @@ 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. + weight_shape = weight.shape + kernel_shape: Optional[Sequence[int]] = None + if weight_shape is not None and len(weight_shape) > 3: + spatial_shape = weight_shape[2:-1] + if all(isinstance(dim, int) for dim in spatial_shape): + kernel_shape = list(spatial_shape) + if transposed: result_real = op.Sub( op.ConvTranspose( @@ -2398,6 +2432,7 @@ def _aten_convolution_complex_onnx( group=groups, dilations=dilations, output_padding=output_padding, + kernel_shape=kernel_shape, ), op.ConvTranspose( input_imag, @@ -2408,6 +2443,7 @@ def _aten_convolution_complex_onnx( group=groups, dilations=dilations, output_padding=output_padding, + kernel_shape=kernel_shape, ), ) result_imag = op.Add( @@ -2420,6 +2456,7 @@ def _aten_convolution_complex_onnx( group=groups, dilations=dilations, output_padding=output_padding, + kernel_shape=kernel_shape, ), op.ConvTranspose( input_imag, @@ -2430,6 +2467,7 @@ def _aten_convolution_complex_onnx( group=groups, dilations=dilations, output_padding=output_padding, + kernel_shape=kernel_shape, ), ) else: @@ -2442,6 +2480,7 @@ def _aten_convolution_complex_onnx( pads=pads, group=groups, dilations=dilations, + kernel_shape=kernel_shape, ), op.Conv( input_imag, @@ -2451,6 +2490,7 @@ def _aten_convolution_complex_onnx( pads=pads, group=groups, dilations=dilations, + kernel_shape=kernel_shape, ), ) @@ -2463,6 +2503,7 @@ def _aten_convolution_complex_onnx( pads=pads, group=groups, dilations=dilations, + kernel_shape=kernel_shape, ), op.Conv( input_imag, @@ -2472,6 +2513,7 @@ def _aten_convolution_complex_onnx( pads=pads, group=groups, dilations=dilations, + kernel_shape=kernel_shape, ), ) From 2f3bde6cb21a7ba76229b7e417e7670b98d105c3 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Thu, 23 Jul 2026 22:35:52 +0000 Subject: [PATCH 3/5] Add regression tests for Conv kernel_shape attribute presence/omission --- .../function_libs/torch_lib/e2e_ops_tests.py | 51 +++++++++++++++++++ 1 file changed, 51 insertions(+) diff --git a/tests/function_libs/torch_lib/e2e_ops_tests.py b/tests/function_libs/torch_lib/e2e_ops_tests.py index e6b481f2c0..4e056a02eb 100644 --- a/tests/function_libs/torch_lib/e2e_ops_tests.py +++ b/tests/function_libs/torch_lib/e2e_ops_tests.py @@ -1266,6 +1266,57 @@ 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_proto.graph.node if n.op_type == "Conv" + ] + self.assertEqual(len(conv_nodes), 1) + attribute_names = {a.name for a in conv_nodes[0].attribute} + self.assertIn("kernel_shape", attribute_names) + + 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_proto.graph.node if n.op_type == "Conv" + ] + self.assertEqual(len(conv_nodes), 1) + attribute_names = {a.name for a in conv_nodes[0].attribute} + self.assertNotIn("kernel_shape", attribute_names) + @unittest.skip("see https://github.com/pytorch/pytorch/issues/174668") def test_aten_histc_float16(self): class Model(torch.nn.Module): From 833c74b96f67b06868ecab7941e73517e464df61 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Thu, 23 Jul 2026 22:43:03 +0000 Subject: [PATCH 4/5] Reuse _conv_kernel_shape for complex path; add ConvTranspose/complex kernel_shape tests --- .../function_libs/torch_lib/ops/core.py | 23 +++--- .../function_libs/torch_lib/e2e_ops_tests.py | 73 ++++++++++++++++++- 2 files changed, 83 insertions(+), 13 deletions(-) diff --git a/onnxscript/function_libs/torch_lib/ops/core.py b/onnxscript/function_libs/torch_lib/ops/core.py index 805ab1ecee..c31c24eefb 100644 --- a/onnxscript/function_libs/torch_lib/ops/core.py +++ b/onnxscript/function_libs/torch_lib/ops/core.py @@ -2305,18 +2305,24 @@ def aten_convolution( return result -def _conv_kernel_shape(weight: TFloat) -> Optional[Sequence[int]]: +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). Returns None when the spatial shape is not fully - static, in which case the attribute is omitted and inferred by ONNX. + (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 or len(shape) <= 2: + 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 - kernel_shape = shape[2:] if any(not isinstance(dim, int) for dim in kernel_shape): return None return list(kernel_shape) @@ -2414,12 +2420,7 @@ def _aten_convolution_complex_onnx( # 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. - weight_shape = weight.shape - kernel_shape: Optional[Sequence[int]] = None - if weight_shape is not None and len(weight_shape) > 3: - spatial_shape = weight_shape[2:-1] - if all(isinstance(dim, int) for dim in spatial_shape): - kernel_shape = list(spatial_shape) + kernel_shape = _conv_kernel_shape(weight, complex=True) if transposed: result_real = op.Sub( diff --git a/tests/function_libs/torch_lib/e2e_ops_tests.py b/tests/function_libs/torch_lib/e2e_ops_tests.py index 4e056a02eb..92a4277e7d 100644 --- a/tests/function_libs/torch_lib/e2e_ops_tests.py +++ b/tests/function_libs/torch_lib/e2e_ops_tests.py @@ -1290,8 +1290,77 @@ def forward(self, x): n for n in onnx_program.model_proto.graph.node if n.op_type == "Conv" ] self.assertEqual(len(conv_nodes), 1) - attribute_names = {a.name for a in conv_nodes[0].attribute} - self.assertIn("kernel_shape", attribute_names) + kernel_shape_attrs = [ + a for a in conv_nodes[0].attribute if a.name == "kernel_shape" + ] + self.assertEqual(len(kernel_shape_attrs), 1) + self.assertEqual(list(kernel_shape_attrs[0].ints), [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_proto.graph.node + if n.op_type == "ConvTranspose" + ] + self.assertEqual(len(conv_nodes), 1) + kernel_shape_attrs = [ + a for a in conv_nodes[0].attribute if a.name == "kernel_shape" + ] + self.assertEqual(len(kernel_shape_attrs), 1) + self.assertEqual(list(kernel_shape_attrs[0].ints), [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_proto.graph.node if n.op_type == "Conv" + ] + self.assertGreater(len(conv_nodes), 0) + for node in conv_nodes: + kernel_shape_attrs = [ + a for a in node.attribute if a.name == "kernel_shape" + ] + self.assertEqual(len(kernel_shape_attrs), 1) + self.assertEqual(list(kernel_shape_attrs[0].ints), [3, 3]) def test_aten_convolution_omits_kernel_shape_when_dynamic(self): # Regression test for https://github.com/microsoft/onnxscript/pull/2972 From eacee056e50ebf101b8c7b75afaff03f648f522f Mon Sep 17 00:00:00 2001 From: titaiwangms Date: Fri, 24 Jul 2026 18:50:03 +0000 Subject: [PATCH 5/5] use ir-py api --- .../function_libs/torch_lib/e2e_ops_tests.py | 55 +++++++------------ 1 file changed, 19 insertions(+), 36 deletions(-) diff --git a/tests/function_libs/torch_lib/e2e_ops_tests.py b/tests/function_libs/torch_lib/e2e_ops_tests.py index 92a4277e7d..b07f57ce20 100644 --- a/tests/function_libs/torch_lib/e2e_ops_tests.py +++ b/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 @@ -1286,15 +1286,11 @@ def forward(self, x): ) _testing.assert_onnx_program(onnx_program) - conv_nodes = [ - n for n in onnx_program.model_proto.graph.node if n.op_type == "Conv" - ] + conv_nodes = [n for n in onnx_program.model.graph if n.op_type == "Conv"] self.assertEqual(len(conv_nodes), 1) - kernel_shape_attrs = [ - a for a in conv_nodes[0].attribute if a.name == "kernel_shape" - ] + kernel_shape_attrs = [conv_nodes[0].attributes["kernel_shape"]] self.assertEqual(len(kernel_shape_attrs), 1) - self.assertEqual(list(kernel_shape_attrs[0].ints), [3, 3]) + 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 @@ -1316,17 +1312,11 @@ def forward(self, x): ) _testing.assert_onnx_program(onnx_program) - conv_nodes = [ - n - for n in onnx_program.model_proto.graph.node - if n.op_type == "ConvTranspose" - ] + conv_nodes = [n for n in onnx_program.model.graph if n.op_type == "ConvTranspose"] self.assertEqual(len(conv_nodes), 1) - kernel_shape_attrs = [ - a for a in conv_nodes[0].attribute if a.name == "kernel_shape" - ] - self.assertEqual(len(kernel_shape_attrs), 1) - self.assertEqual(list(kernel_shape_attrs[0].ints), [3, 3]) + 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 @@ -1351,16 +1341,12 @@ def forward(self, x): ) _testing.assert_onnx_program(onnx_program) - conv_nodes = [ - n for n in onnx_program.model_proto.graph.node if n.op_type == "Conv" - ] + 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: - kernel_shape_attrs = [ - a for a in node.attribute if a.name == "kernel_shape" - ] - self.assertEqual(len(kernel_shape_attrs), 1) - self.assertEqual(list(kernel_shape_attrs[0].ints), [3, 3]) + 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 @@ -1379,12 +1365,9 @@ def forward(self, x, weight): ) _testing.assert_onnx_program(onnx_program) - conv_nodes = [ - n for n in onnx_program.model_proto.graph.node if n.op_type == "Conv" - ] + conv_nodes = [n for n in onnx_program.model.graph if n.op_type == "Conv"] self.assertEqual(len(conv_nodes), 1) - attribute_names = {a.name for a in conv_nodes[0].attribute} - self.assertNotIn("kernel_shape", attribute_names) + self.assertNotIn("kernel_shape", conv_nodes[0].attributes) @unittest.skip("see https://github.com/pytorch/pytorch/issues/174668") def test_aten_histc_float16(self):