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.64%. Comparing base (69a6938) to head (1501909).
✅ All tests successful. No failed tests found.

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.65%   72.64%   -0.01%     
==========================================
  Files         259      259              
  Lines       31766    31776      +10     
  Branches     3007     3010       +3     
==========================================
+ Hits        23078    23084       +6     
- Misses       7670     7672       +2     
- Partials     1018     1020       +2

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

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

+
+    # If offs is None, it uses the "dense" / "batch" mode where groups are implicit in the batch dimension.
+    # self: (G, M, K), mat2: (G, K, N) -> (G, M, N)
+    # TODO: Implement sparse mode when offs is not None.


Could you raise a not implemented error?

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

+
+def _mock_grouped_mm(self, mat2, offs=None, bias=None, out_dtype=None):
+    res = torch.matmul(self, mat2)
+    if bias is not None:
+        res = res + bias


+def _mock_grouped_mm(self, mat2, offs=None, bias=None, out_dtype=None):
+    res = torch.matmul(self, mat2)
+    if bias is not None:
+        res = res + bias


+
+    # If offs is None, it uses the "dense" / "batch" mode where groups are implicit in the batch dimension.
+    # self: (G, M, K), mat2: (G, K, N) -> (G, M, N)
+    # TODO: Implement sparse mode when offs is not None.


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

+        yield opinfo_core.SampleInput(self_t, args=(mat2_t,))
+
+        # Test with bias
+        g, m, k = self_shape


+        yield opinfo_core.SampleInput(self_t, args=(mat2_t,))
+
+        # Test with bias
+        g, m, k = self_shape


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

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

Sid-V5 and others added 4 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>

Address review comments on grouped_mm

1501909

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

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Add converter for `aten::_grouped_mm.default`#2805

Add converter for `aten::_grouped_mm.default`#2805
Sid-V5 wants to merge 4 commits into
microsoft:mainfrom
Sid-V5:add-grouped-mm-2795

Sid-V5 commented Feb 12, 2026 •

edited

Loading

Uh oh!

Sid-V5 commented Feb 12, 2026

Uh oh!

codecov Bot commented Feb 16, 2026 •

edited

Loading

Uh oh!

justinchuby Mar 13, 2026

Uh oh!

justinchuby Mar 13, 2026

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

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!

Check warning

Check warning

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

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!

justinchuby Mar 13, 2026

Choose a reason for hiding this comment

Uh oh!

justinchuby Mar 13, 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!

justinchuby commented Jun 19, 2026

Uh oh!

justinchuby commented Jul 8, 2026

Uh oh!

Sid-V5 commented Jul 8, 2026

Uh oh!

Check warning

Check warning

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

Sid-V5 commented Feb 12, 2026 •

edited

Loading

codecov Bot commented Feb 16, 2026 •

edited

Loading