Add converter for `aten::_grouped_mm.default` by Sid-V5 · Pull Request #2805 · microsoft/onnxscript

Sid (Sid-V5) · 2026-02-12T20:50:54Z

Implemented the converter for aten::_grouped_mm.default to address #2795.

Changes

Added aten_grouped_mm function in onnxscript/function_libs/torch_lib/ops/core.py

Implementation Details

The aten::_grouped_mm operator performs grouped matrix multiplication. This implementation handles the batch/dense mode (when offs is None), where groups are implicit in the batch dimension:

self: (G, M, K), mat2: (G, K, N) → result: (G, M, N)
Uses op.MatMul for the core computation
Supports optional bias addition via op.Add
Supports optional out_dtype casting via op.Cast

The offset-based mode (when offs is provided) raises NotImplementedError, as it requires segment-level matrix multiplications that are not directly expressible with standard ONNX operators.

Testing

The function follows the same patterns as other converters in core.py (e.g., aten_bmm, aten_mm) and uses the @torch_op decorator for automatic registration.

Fixes #2795

Sid (Sid-V5) · 2026-02-12T20:54:38Z

@microsoft-github-policy-service agree

codecov · 2026-02-16T18:23:52Z

Codecov Report

❌ Patch coverage is 60.00000% with 4 lines in your changes missing coverage. Please review.
✅ Project coverage is 72.64%. Comparing base (3ba2bf7) to head (9ba30a4).

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

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2805      +/-   ##
==========================================
- Coverage   72.64%   72.64%   -0.01%     
==========================================
  Files         265      265              
  Lines       32251    32261      +10     
  Branches     3050     3053       +3     
==========================================
+ Hits        23429    23435       +6     
- Misses       7786     7788       +2     
- Partials     1036     1038       +2

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

Copilot

Pull request overview

Adds an ONNXScript converter for PyTorch’s aten::_grouped_mm (dense/batched mode) and wires it into the TorchLib op test infrastructure to address missing dispatch for aten._grouped_mm.default (#2795).

Changes:

Implemented aten_grouped_mm converter using MatMul, optional Add (bias), and optional Cast (out_dtype).
Registered the new op in TorchLib’s tested op list.
Added OpInfo + sample inputs scaffolding for _grouped_mm in the extra op database.

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 4 comments.

File	Description
onnxscript/function_libs/torch_lib/ops/core.py	Adds the `aten::_grouped_mm` converter implementation.
tests/function_libs/torch_lib/ops_test_data.py	Enables testing by adding a `TorchLibOpInfo` entry for `ops.aten._grouped_mm`.
tests/function_libs/torch_lib/extra_opinfo.py	Adds OpInfo registration and sample input generation for `_grouped_mm`.

Justin Chu (justinchuby) · 2026-06-19T17:59:40Z

Two items. (1) The CONFLICTING flag is trivial — a git merge-tree against main resolves cleanly (insertions into non-overlapping regions), so a rebase clears it. (2) Real correctness issue: the offs argument is silently ignored — the code has only a # TODO and falls through to MatMul (~core.py L4515-4523), but the PR description says offset mode "raises NotImplementedError." With offs provided (the ragged/MoE case) this emits silently-incorrect results instead of failing fast; please add an explicit if offs is not None: raise NotImplementedError(...). Also, the OpInfo currently compares against a hand-written torch.matmul-based mock rather than the real torch._grouped_mm, so it's near-tautological — worth strengthening toward real-op parity and covering bias / out_dtype.

Justin Chu (justinchuby) · 2026-07-08T16:41:05Z

Sid (@Sid-V5) would you like to update this PR?

Implements the converter for aten::_grouped_mm.default to address issue microsoft#2795. Handles the batch/dense mode where groups are implicit in the batch dimension using MatMul, with optional bias addition and dtype casting.

…t cases

Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>

Sid (Sid-V5) · 2026-07-08T17:43:48Z

Justin Chu (@justinchuby)
Rebased onto main to resolve the conflict.
updated core.py to raise NotImplementedError if offs is not None, and updated the OpInfo tests to cover bias/out_dtype and use the real torch._grouped_mm when available

Copilot

Pull request overview

Copilot reviewed 3 out of 3 changed files in this pull request and generated 4 comments.

- Changed out_dtype default from Optional[int]=None to int=-1, matching the sentinel convention used throughout core.py (e.g. aten_sum). Guards for both None and -1 to prevent invalid op.Cast(..., to=-1). - Pass offs, bias, out_dtype as keyword arguments in SampleInput and in the native torch.ops.aten._grouped_mm call, since they are keyword-only in the aten schema (after *). This prevents the try: path from silently failing with TypeError. Addresses Copilot review comments and RUFF/RUF059 lint feedback.

Justin Chu (justinchuby) · 2026-07-23T14:35:13Z

Could you fix lint?

Copilot

Pull request overview

Copilot reviewed 3 out of 3 changed files in this pull request and generated 1 comment.

Copilot

Pull request overview

Copilot reviewed 3 out of 3 changed files in this pull request and generated no new comments.

Copilot

Pull request overview

Copilot reviewed 3 out of 3 changed files in this pull request and generated no new comments.

Copilot

Pull request overview

Copilot reviewed 3 out of 3 changed files in this pull request and generated no new comments.

Suppressed comments (1)

onnxscript/function_libs/torch_lib/ops/core.py:4603

out_dtype handling should follow the existing pattern used elsewhere (e.g., quantized_decomposed.py:62-66): treat -1 and None as “not set”, and assert that any provided dtype id is valid (>0) before passing it to op.Cast. This avoids accidentally emitting an invalid ONNX Cast(to=0) if a caller ever forwards an unset/unknown dtype value.

    if out_dtype is not None and out_dtype != -1:
        res = op.Cast(res, to=out_dtype)
    return res

Justin Chu (justinchuby) · 2026-08-01T02:54:06Z

Sid (@Sid-V5) just my inline reviews and it should be ready to merge after that

Justin Chu (justinchuby) · 2026-08-31T20:24:51Z

+    opinfo_core.OpInfo(
+        "ops.aten._grouped_mm",
+        aten_name="_grouped_mm",
+        op=_mock_grouped_mm,


As commented, please change this to

Suggested change

op=_mock_grouped_mm,

op=torch.ops.aten._grouped_mm,

Then in TorchLibOpInfo, you may skip the test if not hasattr(torch.ops.aten, "_grouped_mm")

Justin Chu (justinchuby)

Nit as commented

… unavailable - Update OpInfo in extra_opinfo.py to reference torch.ops.aten._grouped_mm directly instead of _mock_grouped_mm. - Remove _mock_grouped_mm reference helper. - Add .skip(enabled_if=not hasattr(torch.ops.aten, '_grouped_mm')) to TorchLibOpInfo in ops_test_data.py to cleanly skip on PyTorch versions without _grouped_mm support.

github-project-automation Bot added this to ONNX Script Review Board Feb 12, 2026

github-project-automation Bot moved this to Todo in ONNX Script Review Board Feb 12, 2026

Justin Chu (justinchuby) assigned Xavier Dupré (xadupre) Feb 13, 2026

Justin Chu (justinchuby) reviewed Mar 13, 2026

View reviewed changes

Comment thread onnxscript/function_libs/torch_lib/ops/core.py Outdated

Justin Chu (justinchuby) reviewed Mar 13, 2026

View reviewed changes

Comment thread tests/function_libs/torch_lib/extra_opinfo.py Outdated

Justin Chu (justinchuby) added the module: torchlib Related to the torch/aten function lib in development label Mar 13, 2026

Justin Chu (justinchuby) requested a review from Copilot March 13, 2026 16:26

Copilot started reviewing on behalf of Justin Chu (justinchuby) March 13, 2026 16:27 View session

github-advanced-security AI found potential problems Mar 13, 2026

View reviewed changes

Comment thread tests/function_libs/torch_lib/extra_opinfo.py Fixed

Comment thread tests/function_libs/torch_lib/extra_opinfo.py Fixed

Copilot AI reviewed Mar 13, 2026

View reviewed changes

Comment thread tests/function_libs/torch_lib/extra_opinfo.py Outdated

Comment thread tests/function_libs/torch_lib/extra_opinfo.py Outdated

Comment thread onnxscript/function_libs/torch_lib/ops/core.py Outdated

Comment thread tests/function_libs/torch_lib/extra_opinfo.py

Sid (Sid-V5) and others added 3 commits July 8, 2026 23:03

Add aten::_grouped_mm converter implementation

dc1281d

Implements the converter for aten::_grouped_mm.default to address issue microsoft#2795. Handles the batch/dense mode where groups are implicit in the batch dimension using MatMul, with optional bias addition and dtype casting.

Fix grouped_mm: Handle optional arguments correctly and add basic tes…

d61ccf7

…t cases

Potential fix for pull request finding

df7fe91

Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>

Sid (Sid-V5) force-pushed the add-grouped-mm-2795 branch from d73b04f to 1501909 Compare July 8, 2026 17:42

github-advanced-security AI found potential problems Jul 8, 2026

View reviewed changes

Comment thread tests/function_libs/torch_lib/extra_opinfo.py Fixed

Comment thread tests/function_libs/torch_lib/extra_opinfo.py Fixed

Address review comments on grouped_mm

0fa4617

Sid (Sid-V5) force-pushed the add-grouped-mm-2795 branch from 1501909 to 0fa4617 Compare July 9, 2026 08:13

Merge branch 'main' into add-grouped-mm-2795

3940fa8

G. Ramalingam (gramalingam) requested a review from Copilot July 15, 2026 23:50

Copilot started reviewing on behalf of G. Ramalingam (gramalingam) July 15, 2026 23:50 View session

Copilot AI reviewed Jul 15, 2026

View reviewed changes

Comment thread onnxscript/function_libs/torch_lib/ops/core.py

Comment thread tests/function_libs/torch_lib/extra_opinfo.py

Comment thread tests/function_libs/torch_lib/extra_opinfo.py

Comment thread tests/function_libs/torch_lib/extra_opinfo.py Outdated

Sid (Sid-V5) and others added 2 commits July 16, 2026 10:38

Merge branch 'main' into add-grouped-mm-2795

74bac5e

G. Ramalingam (gramalingam) approved these changes Jul 16, 2026

View reviewed changes

fix lint formatting

dc64093

Copilot AI review requested due to automatic review settings July 23, 2026 16:57

Copilot started reviewing on behalf of Sid (Sid-V5) July 23, 2026 16:58 View session

Copilot AI reviewed Jul 23, 2026

View reviewed changes

Comment thread tests/function_libs/torch_lib/extra_opinfo.py Outdated

narrow exception catch to TypeError and RuntimeError

79ebf71

Copilot AI review requested due to automatic review settings July 23, 2026 17:07

Copilot started reviewing on behalf of Sid (Sid-V5) July 23, 2026 17:08 View session

Copilot AI reviewed Jul 23, 2026

View reviewed changes

Merge branch 'main' into add-grouped-mm-2795

0cef08e

Copilot AI review requested due to automatic review settings July 28, 2026 12:32

Copilot started reviewing on behalf of Sid (Sid-V5) July 28, 2026 12:32 View session

Copilot AI reviewed Jul 28, 2026

View reviewed changes

Merge branch 'main' into add-grouped-mm-2795

6483df0

Copilot AI review requested due to automatic review settings July 31, 2026 17:34

Copilot started reviewing on behalf of Sid (Sid-V5) July 31, 2026 17:35 View session

Copilot AI reviewed Jul 31, 2026

View reviewed changes

Sid (Sid-V5) added 2 commits August 20, 2026 09:21

Merge branch 'main' into add-grouped-mm-2795

3306a31

Merge branch 'main' into add-grouped-mm-2795

9ba30a4

Justin Chu (justinchuby) approved these changes Aug 31, 2026

View reviewed changes

github-project-automation Bot moved this from Todo to Done in ONNX Script Review Board Aug 31, 2026

Justin Chu (justinchuby) enabled auto-merge (squash) August 31, 2026 20:22

Justin Chu (justinchuby) disabled auto-merge August 31, 2026 20:23

Justin Chu (justinchuby) reviewed Aug 31, 2026

View reviewed changes

Justin Chu (justinchuby) requested changes Aug 31, 2026

View reviewed changes

github-project-automation Bot moved this from Done to In Progress in ONNX Script Review Board Aug 31, 2026

Uh oh!

Conversation

Sid (Sid-V5) commented Feb 12, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Changes

Implementation Details

Testing

Uh oh!

Sid (Sid-V5) commented Feb 12, 2026

Uh oh!

codecov Bot commented Feb 16, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Justin Chu (justinchuby) commented Jun 19, 2026

Uh oh!

Justin Chu (justinchuby) commented Jul 8, 2026

Uh oh!

Sid (Sid-V5) commented Jul 8, 2026

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Justin Chu (justinchuby) commented Jul 23, 2026

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Justin Chu (justinchuby) commented Aug 1, 2026

Uh oh!

Justin Chu (justinchuby) Aug 31, 2026

Choose a reason for hiding this comment

Uh oh!

Justin Chu (justinchuby) left a comment

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

Sid (Sid-V5) commented Feb 12, 2026 •

edited

Loading

codecov Bot commented Feb 16, 2026 •

edited

Loading