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

❌ 4 Tests Failed:

Tests completed	Failed	Passed	Skipped
4447	4	4443	949

View the full list of 4 ❄️ flaky test(s)

tests.function_libs.torch_lib.ops_test.TestOutputConsistencyFullGraphCPU::test_output_match_opinfo__arange_start_step_cpu_int32

Flake rate in main: 11.76% (Passed 30 times, Failed 4 times)

Stack Traces | 0.425s run time

.../function_libs/torch_lib/ops_test.py:243: in run_test_output_match
    torch.testing.assert_close(
E   AssertionError: Tensor-likes are not equal!
E   
E   Mismatched elements: 9 / 10 (90.0%)
E   Greatest absolute difference: 5 at index (9,)
E   Greatest relative difference: 1.0 at index (9,)

tests.function_libs.torch_lib.ops_test.TestOutputConsistencyFullGraphCPU::test_output_match_opinfo__arange_start_step_cpu_int64

Flake rate in main: 11.76% (Passed 30 times, Failed 4 times)

Stack Traces | 0.554s run time

.../function_libs/torch_lib/ops_test_common.py:593: in _capture_graph_and_evaluate_torch_script_evaluator
    return _ort_session_run(model_proto.SerializeToString(), ort_inputs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
.../function_libs/torch_lib/ops_test_common.py:349: in _ort_session_run
    return session.run(None, ort_inputs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
..../test_torch_nightly/lib/python3.11.../onnxruntime/capi/onnxruntime_inference_collection.py:275: in run
    return self._sess.run(output_names, input_feed, run_options)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
E   onnxruntime.capi.onnxruntime_pybind11_state.InvalidArgument: [ONNXRuntimeError] : 2 : INVALID_ARGUMENT : Non-zero status code returned while running Range node. Name:'node_Range_6' Status Message: delta in Range operator can not be zero!

The above exception was the direct cause of the following exception:
.../function_libs/torch_lib/ops_test.py:206: in run_test_output_match
    function_output = function_executor(test_name, reference_torch_outputs)(
.../function_libs/torch_lib/ops_test_common.py:607: in _capture_graph_and_evaluate_torch_script_evaluator
    raise RuntimeError(
E   RuntimeError: ONNX Runtime failed to evaluate:
E   Inputs:
E   {}
E   Model:
E   <
E      ir_version: 10,
E      opset_import: ["" : 18, "pkg.torch.onnx" : 1, "pkg.onnxscript.torch_lib.common" : 1, "pkg.onnxscript.torch_lib" : 1],
E      producer_name: "torch_test"
E   >
E   main_graph () => (int64[10] val_6) 
E      <int64 val_0, int64 val_1, int64 val_2, int64 val_3, float val_4, int64 val_5>
E   {
E      [node_Constant_0] val_0 = Constant <value: tensor = int64 {0}> ()
E      [node_Cast_1] val_1 = Cast <to: int = 7> (val_0)
E      [node_Constant_2] val_2 = Constant <value: tensor = int64 {5}> ()
E      [node_Cast_3] val_3 = Cast <to: int = 7> (val_2)
E      [node_Constant_4] val_4 = Constant <value: tensor = float {-0.5}> ()
E      [node_Cast_5] val_5 = Cast <to: int = 7> (val_4)
E      [node_Range_6] val_6 = Range (val_3, val_1, val_5)
E   }

tests.function_libs.torch_lib.ops_test.TestOutputConsistencyFullGraphCPU::test_output_match_opinfo__logit_cpu_float16

Flake rate in main: 25.00% (Passed 3 times, Failed 1 times)

Stack Traces | 0.154s run time

.../function_libs/torch_lib/ops_test.py:243: in run_test_output_match
    torch.testing.assert_close(
E   AssertionError: Tensor-likes are not close!
E   
E   Mismatched elements: 42 / 125 (33.6%)
E   Greatest absolute difference: 0.8115234375 at index (0, 0, 1) (up to 0.0007 allowed)
E   Greatest relative difference: 2.0 at index (0, 0, 1) (up to 0.1 allowed)

tests.function_libs.torch_lib.ops_test.TestOutputConsistencyFullGraphCPU::test_output_match_opinfo__logit_cpu_float32

Flake rate in main: 25.00% (Passed 3 times, Failed 1 times)

Stack Traces | 0.146s run time

.../function_libs/torch_lib/ops_test.py:243: in run_test_output_match
    torch.testing.assert_close(
E   AssertionError: Tensor-likes are not close!
E   
E   Mismatched elements: 44 / 125 (35.2%)
E   Greatest absolute difference: 0.8109303712844849 at index (0, 0, 1) (up to 1e-05 allowed)
E   Greatest relative difference: 2.000000238418579 at index (0, 0, 1) (up to 1.3e-06 allowed)

To view more test analytics, go to the Test Analytics Dashboard
_{📋 Got 3 mins? Take this short survey to help us improve Test Analytics.}

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

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.

+            return torch.ops.aten._grouped_mm(
+                self, mat2, offs=offs, bias=bias, out_dtype=out_dtype
+            )
+        except Exception:


Copilot

Pull request overview

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

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

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

return torch.ops.aten._grouped_mm(

self, mat2, offs=offs, bias=bias, out_dtype=out_dtype

)

except Exception:

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

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!

❌ 4 Tests Failed:

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!

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!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

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