Add converter for aten::_grouped_mm.default

[Sid-V5]

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]

@microsoft-github-policy-service agree

[codecov]

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

Could you raise a not implemented error?

[justinchuby]

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]

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]

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

[Sid-V5]

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