Set kernel_shape attribute on Conv/ConvTranspose in torchlib

[Copilot]

Per the ONNX Conv spec, kernel_shape describes the spatial shape of the convolution kernel, but torchlib omitted it, leaving the attribute absent from exported models.

Changes

• _conv_kernel_shape helper: derives kernel_shape from the weight's spatial dims (all dims except the leading output/input channel dims). Returns None when the spatial shape is not fully static, so the attribute is omitted and left to ONNX inference rather than emitting an incorrect value.
• _aten_convolution_onnx: passes kernel_shape to op.Conv and op.ConvTranspose (covers aten_conv1d/conv2d/conv3d/convolution and the transpose path).
• _aten_convolution_complex_onnx: computes kernel_shape from the complex weight, excluding the trailing size-2 real/imag dimension, and applies it to all Conv/ConvTranspose calls.

Result

# static weight (4, 3, 3, 3)
Conv {..., 'strides': (1, 1), 'dilations': (1, 1), 'kernel_shape': (3, 3)}

# dynamic spatial weight (N, 3, k, k) -> attribute omitted
Conv {..., 'strides': (1, 1), 'dilations': (1, 1)}

• Fixes torchlib does not specify kernel_shape attribute for op.Conv #2971

[titaiwangms]

Fix pytorch/pytorch#190890

[Copilot]

Pull request overview

This PR aligns torchlib’s aten::convolution ONNX lowering with the ONNX Conv/ConvTranspose specification by explicitly setting the kernel_shape attribute when it can be derived reliably from the (static) convolution weight tensor shape, and omitting it when the spatial kernel shape is not fully static.

Changes:

• Added _conv_kernel_shape helper to derive kernel_shape from weight.shape[2:] when all spatial dims are statically known.
• Updated _aten_convolution_onnx to pass kernel_shape through to op.Conv and op.ConvTranspose.
• Updated _aten_convolution_complex_onnx to compute kernel_shape for complex weights (excluding the trailing real/imag dimension) and pass it to all Conv/ConvTranspose calls.

[codecov]

Codecov Report

❌ Patch coverage is 50.00000% with 6 lines in your changes missing coverage. Please review.
✅ Project coverage is 72.63%. Comparing base (00de62e) to head (833c74b).
✅ All tests successful. No failed tests found.

Files with missing lines	Patch %	Lines
onnxscript/function_libs/torch_lib/ops/core.py	50.00%	3 Missing and 3 partials ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2972      +/-   ##
==========================================
- Coverage   72.64%   72.63%   -0.01%     
==========================================
  Files         265      265              
  Lines       32192    32204      +12     
  Branches     3038     3041       +3     
==========================================
+ Hits        23385    23391       +6     
- Misses       7776     7779       +3     
- Partials     1031     1034       +3     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

[titaiwangms]

Review summary (spun up review team: code-reviewer, critical-reviewer, integration-reviewer, qa-tester)

Verdict: fix is correct and verified to resolve pytorch/pytorch#190890.

QA verification

Ran the exact repro from pytorch/pytorch#190890 against this branch:

• With fix: exported Conv node has kernel_shape: [3, 3] ✅
• With core.py reverted to upstream/main (isolating this diff): kernel_shape absent ❌, confirming this change is what fixes it.
• Also tried dynamo=True export with dynamic H/W input shapes — kernel_shape=[3, 3] still correctly emitted (it's derived from the static weight, not the dynamic input), no crash.
• Not exercised: the complex-conv path, ConvTranspose, and the "kernel_shape omitted for symbolic weight shape" branch of _conv_kernel_shape.

Code review findings

No Critical/Major bugs found. Confirmed:

• kernel_shape=None is safely filtered out by the attribute converter (no invalid/empty attribute emitted).
• weight.shape[2:] correctly extracts spatial dims regardless of groups/depthwise conv.
• All Conv/ConvTranspose call sites in torch_lib (conv1d/2d/3d, transposed, complex) route through the two patched functions (_aten_convolution_onnx, _aten_convolution_complex_onnx) — nothing missed.
• Existing rewriter passes (_fuse_pad_into_conv.py, _fuse_batchnorm.py, etc.) already handle/preserve kernel_shape safely; _fuse_pad_into_conv.py even benefits from it being set explicitly instead of inferring from weight itself.

Minor (raised independently by two reviewers): No regression test was added asserting kernel_shape is present/absent in the exported graph attributes. Existing torch_lib tests only check numeric output equality (ONNX Runtime infers the attribute when absent), so this wouldn't catch a future regression of #2971. Suggest adding a graph-attribute test covering: static weight → kernel_shape present with correct value, and symbolic/non-static weight spatial shape → attribute omitted (also for the complex-conv path and ConvTranspose).

Minor: The complex-conv path (_aten_convolution_complex_onnx) duplicates the shape-extraction logic inline instead of reusing _conv_kernel_shape, which risks drift if the extraction rule changes later.

No blocking issues; the missing test coverage is worth addressing before/after merge but doesn't affect correctness of the fix as verified above.

[titaiwangms]

Review summary (spun up review team: code-reviewer, critical-reviewer, integration-reviewer, qa-tester)

Verdict: fix is correct and verified to resolve pytorch/pytorch#190890.

QA verification

Ran the exact repro from pytorch/pytorch#190890 against this branch:

• With fix: exported Conv node has kernel_shape: [3, 3] ✅
• With core.py reverted to upstream/main (isolating this diff): kernel_shape absent ❌, confirming this change is what fixes it.
• Also tried dynamo=True export with dynamic H/W input shapes — kernel_shape=[3, 3] still correctly emitted (it's derived from the static weight, not the dynamic input), no crash.
• Not exercised: the complex-conv path, ConvTranspose, and the "kernel_shape omitted for symbolic weight shape" branch of _conv_kernel_shape.

Code review findings

No Critical/Major bugs found. Confirmed:

• kernel_shape=None is safely filtered out by the attribute converter (no invalid/empty attribute emitted).
• weight.shape[2:] correctly extracts spatial dims regardless of groups/depthwise conv.
• All Conv/ConvTranspose call sites in torch_lib (conv1d/2d/3d, transposed, complex) route through the two patched functions (_aten_convolution_onnx, _aten_convolution_complex_onnx) — nothing missed.
• Existing rewriter passes (_fuse_pad_into_conv.py, _fuse_batchnorm.py, etc.) already handle/preserve kernel_shape safely; _fuse_pad_into_conv.py even benefits from it being set explicitly instead of inferring from weight itself.

Minor (raised independently by two reviewers): No regression test was added asserting kernel_shape is present/absent in the exported graph attributes. Existing torch_lib tests only check numeric output equality (ONNX Runtime infers the attribute when absent), so this wouldn't catch a future regression of #2971. Suggest adding a graph-attribute test covering: static weight → kernel_shape present with correct value, and symbolic/non-static weight spatial shape → attribute omitted (also for the complex-conv path and ConvTranspose).

Minor: The complex-conv path (_aten_convolution_complex_onnx) duplicates the shape-extraction logic inline instead of reusing _conv_kernel_shape, which risks drift if the extraction rule changes later.

No blocking issues; the missing test coverage is worth addressing before/after merge but doesn't affect correctness of the fix as verified above.

@copilot address these

[Copilot]

Review summary (spun up review team: code-reviewer, critical-reviewer, integration-reviewer, qa-tester)

**Verdict: fix is correct and verified to resolve [pytorch/pytorch#190890](https://githu...

Addressed both minor points:

• Reused _conv_kernel_shape in _aten_convolution_complex_onnx via a new complex flag, removing the duplicated inline shape-extraction logic.
• Extended the graph-attribute regression tests to cover ConvTranspose and the complex-conv path, and now assert the kernel_shape value ([3, 3]) in addition to presence/omission for Conv.

Existing convolution numerical opinfo tests still pass.

[justinchuby]

Suggested change

	n for n in onnx_program.model_proto.graph.node if n.op_type == "Conv"
	n for n in onnx_program.model.graph if n.op_type == "Conv"

Use the IR object

[justinchuby]

Ditto