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

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-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.62%. Comparing base (9ef0aac) to head (0cef08e).

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.63%   72.62%   -0.01%     
==========================================
  Files         265      265              
  Lines       32204    32214      +10     
  Branches     3041     3044       +3     
==========================================
+ Hits        23391    23397       +6     
- Misses       7779     7781       +2     
- Partials     1034     1036       +2

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

justinchuby · 2026-03-13T16:24:30Z

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


Why not use

Suggested change

op=_mock_grouped_mm,

op=torch.ops.aten._grouped_mm,

?

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`.

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.

justinchuby · 2026-07-08T16:41:05Z

@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-V5 · 2026-07-08T17:43:48Z

@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.

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

justinchuby · 2026-08-01T02:54:06Z

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

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

justinchuby assigned xadupre Feb 13, 2026

justinchuby reviewed Mar 13, 2026

View reviewed changes

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

justinchuby reviewed Mar 13, 2026

View reviewed changes

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

justinchuby requested a review from Copilot March 13, 2026 16:26

Copilot started reviewing on behalf of 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

Comment thread tests/function_libs/torch_lib/extra_opinfo.py

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

Comment thread tests/function_libs/torch_lib/extra_opinfo.py

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-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-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

gramalingam requested a review from Copilot July 15, 2026 23:50

Copilot started reviewing on behalf of 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

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

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

74bac5e

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-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-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-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-V5 July 31, 2026 17:35 View session

Copilot AI reviewed Jul 31, 2026

View reviewed changes

Uh oh!

Conversation

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-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!

justinchuby Mar 13, 2026

Choose a reason for hiding this comment

Uh oh!

justinchuby Jul 23, 2026

Choose a reason for hiding this comment

Uh oh!

justinchuby Jul 23, 2026

Choose a reason for hiding this comment

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!

justinchuby commented Jun 19, 2026

Uh oh!

justinchuby commented Jul 8, 2026

Uh oh!

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!

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!

justinchuby commented Aug 1, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

Sid-V5 commented Feb 12, 2026 •

edited

Loading

codecov Bot commented Feb 16, 2026 •

edited

Loading