From dc1281d01a2f33dafb33ce345c1242a9308b6321 Mon Sep 17 00:00:00 2001 From: Legendx4060 Date: Fri, 13 Feb 2026 02:20:07 +0530 Subject: [PATCH 1/4] Add aten::_grouped_mm converter implementation Implements the converter for aten::_grouped_mm.default to address issue #2795. Handles the batch/dense mode where groups are implicit in the batch dimension using MatMul, with optional bias addition and dtype casting. --- .../function_libs/torch_lib/ops/core.py | 23 +++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/onnxscript/function_libs/torch_lib/ops/core.py b/onnxscript/function_libs/torch_lib/ops/core.py index 418e375c1b..8df0d9e191 100644 --- a/onnxscript/function_libs/torch_lib/ops/core.py +++ b/onnxscript/function_libs/torch_lib/ops/core.py @@ -4510,6 +4510,29 @@ def aten_grid_sampler_3d_backward( raise NotImplementedError() +@torch_op("aten::_grouped_mm") +def aten_grouped_mm( + self: TFloat, + mat2: TFloat, + offs: Optional[TInt] = None, + bias: Optional[TFloat] = None, + out_dtype: Optional[int] = None, +) -> TFloat: + """_grouped_mm(Tensor self, Tensor mat2, *, Tensor? offs=None, Tensor? bias=None, int? out_dtype=None) -> Tensor""" + + # 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) + if offs is None: + res = op.MatMul(self, mat2) + if bias is not None: + res = op.Add(res, bias) + if out_dtype is not None: + res = op.Cast(res, to=out_dtype) + return res + + raise NotImplementedError("aten::_grouped_mm with 'offs' is not supported.") + + def aten_gru_cell( input: TensorType, hx: TensorType, From d61ccf7c5180630ebb4474c4162e1faaa23a74b3 Mon Sep 17 00:00:00 2001 From: Legendx4060 Date: Thu, 12 Mar 2026 18:16:49 +0530 Subject: [PATCH 2/4] Fix grouped_mm: Handle optional arguments correctly and add basic test cases --- .../function_libs/torch_lib/ops/core.py | 18 +++++----- tests/function_libs/torch_lib/extra_opinfo.py | 36 +++++++++++++++++++ .../function_libs/torch_lib/ops_test_data.py | 1 + 3 files changed, 45 insertions(+), 10 deletions(-) diff --git a/onnxscript/function_libs/torch_lib/ops/core.py b/onnxscript/function_libs/torch_lib/ops/core.py index 8df0d9e191..401b51ca4e 100644 --- a/onnxscript/function_libs/torch_lib/ops/core.py +++ b/onnxscript/function_libs/torch_lib/ops/core.py @@ -4510,7 +4510,7 @@ def aten_grid_sampler_3d_backward( raise NotImplementedError() -@torch_op("aten::_grouped_mm") +@torch_op("aten::_grouped_mm", trace_only=True) def aten_grouped_mm( self: TFloat, mat2: TFloat, @@ -4522,15 +4522,13 @@ def aten_grouped_mm( # 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) - if offs is None: - res = op.MatMul(self, mat2) - if bias is not None: - res = op.Add(res, bias) - if out_dtype is not None: - res = op.Cast(res, to=out_dtype) - return res - - raise NotImplementedError("aten::_grouped_mm with 'offs' is not supported.") + # TODO: Implement sparse mode when offs is not None. + res = op.MatMul(self, mat2) + if bias is not None: + res = op.Add(res, bias) + if out_dtype is not None: + res = op.Cast(res, to=out_dtype) + return res def aten_gru_cell( diff --git a/tests/function_libs/torch_lib/extra_opinfo.py b/tests/function_libs/torch_lib/extra_opinfo.py index d925be6877..324af6c23d 100644 --- a/tests/function_libs/torch_lib/extra_opinfo.py +++ b/tests/function_libs/torch_lib/extra_opinfo.py @@ -23,6 +23,34 @@ M = 10 +def sample_inputs_grouped_mm(op_info, device, dtype, requires_grad, **kwargs): + del op_info + del kwargs + + make_arg = functools.partial( + torch_testing.make_tensor, device=device, dtype=dtype, requires_grad=requires_grad + ) + + cases = [ + # (G, M, K), (G, K, N) + ((2, 3, 4), (2, 4, 5)), + ((1, 2, 2), (1, 2, 1)), + ] + + for self_shape, mat2_shape in cases: + self_t = make_arg(self_shape) + mat2_t = make_arg(mat2_shape) + + # Test without bias + yield opinfo_core.SampleInput(self_t, args=(mat2_t,)) + +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 + return res + + def sample_inputs_scalar_tensor(op_info, device, dtype, requires_grad, **kwargs): del op_info del kwargs @@ -3142,4 +3170,12 @@ def sample_inputs_masked_scatter(op_info, device, dtype, requires_grad, **kwargs sample_inputs_func=sample_inputs_masked_scatter, supports_out=False, ), + opinfo_core.OpInfo( + "ops.aten._grouped_mm", + aten_name="_grouped_mm", + op=_mock_grouped_mm, + dtypes=common_dtype.floating_types(), + sample_inputs_func=sample_inputs_grouped_mm, + supports_out=False, + ), ] diff --git a/tests/function_libs/torch_lib/ops_test_data.py b/tests/function_libs/torch_lib/ops_test_data.py index 241ac98cf9..2025c2ae27 100644 --- a/tests/function_libs/torch_lib/ops_test_data.py +++ b/tests/function_libs/torch_lib/ops_test_data.py @@ -750,6 +750,7 @@ def _where_input_wrangler( reason="fixme: ORT does not support empty tensors as input", ), TorchLibOpInfo("ge", core_ops.aten_ge), + TorchLibOpInfo("ops.aten._grouped_mm", core_ops.aten_grouped_mm), TorchLibOpInfo("gt", core_ops.aten_gt), TorchLibOpInfo("histc", core_ops.aten_histc) .skip( From df7fe91d4cb2a2217fe557c275d6bbc79de549c6 Mon Sep 17 00:00:00 2001 From: Sid <138317706+Sid-V5@users.noreply.github.com> Date: Wed, 25 Mar 2026 23:40:07 +0530 Subject: [PATCH 3/4] Potential fix for pull request finding Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --- tests/function_libs/torch_lib/extra_opinfo.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/function_libs/torch_lib/extra_opinfo.py b/tests/function_libs/torch_lib/extra_opinfo.py index 324af6c23d..7033ecb6dd 100644 --- a/tests/function_libs/torch_lib/extra_opinfo.py +++ b/tests/function_libs/torch_lib/extra_opinfo.py @@ -44,6 +44,7 @@ def sample_inputs_grouped_mm(op_info, device, dtype, requires_grad, **kwargs): # Test without bias yield opinfo_core.SampleInput(self_t, args=(mat2_t,)) + def _mock_grouped_mm(self, mat2, offs=None, bias=None, out_dtype=None): res = torch.matmul(self, mat2) if bias is not None: From 0fa461784a4cff069e1bcf16d949b23e031b4f6f Mon Sep 17 00:00:00 2001 From: Legendx4060 Date: Wed, 8 Jul 2026 23:06:50 +0530 Subject: [PATCH 4/4] Address review comments on grouped_mm --- .../function_libs/torch_lib/ops/core.py | 8 +++--- tests/function_libs/torch_lib/extra_opinfo.py | 26 +++++++++++++++++-- 2 files changed, 29 insertions(+), 5 deletions(-) diff --git a/onnxscript/function_libs/torch_lib/ops/core.py b/onnxscript/function_libs/torch_lib/ops/core.py index 401b51ca4e..7e15c069f3 100644 --- a/onnxscript/function_libs/torch_lib/ops/core.py +++ b/onnxscript/function_libs/torch_lib/ops/core.py @@ -4520,9 +4520,11 @@ def aten_grouped_mm( ) -> TFloat: """_grouped_mm(Tensor self, Tensor mat2, *, Tensor? offs=None, Tensor? bias=None, int? out_dtype=None) -> Tensor""" - # 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. + if offs is not None: + raise NotImplementedError( + "Grouped matmul with offsets (ragged/MoE) is not supported." + ) + res = op.MatMul(self, mat2) if bias is not None: res = op.Add(res, bias) diff --git a/tests/function_libs/torch_lib/extra_opinfo.py b/tests/function_libs/torch_lib/extra_opinfo.py index 7033ecb6dd..2c97b73b22 100644 --- a/tests/function_libs/torch_lib/extra_opinfo.py +++ b/tests/function_libs/torch_lib/extra_opinfo.py @@ -28,7 +28,10 @@ def sample_inputs_grouped_mm(op_info, device, dtype, requires_grad, **kwargs): del kwargs make_arg = functools.partial( - torch_testing.make_tensor, device=device, dtype=dtype, requires_grad=requires_grad + torch_testing.make_tensor, + device=device, + dtype=dtype, + requires_grad=requires_grad, ) cases = [ @@ -41,14 +44,33 @@ def sample_inputs_grouped_mm(op_info, device, dtype, requires_grad, **kwargs): self_t = make_arg(self_shape) mat2_t = make_arg(mat2_shape) - # Test without bias + # Test without bias and without out_dtype yield opinfo_core.SampleInput(self_t, args=(mat2_t,)) + # Test with bias + g, _, _ = self_shape + _, _, n = mat2_shape + bias_t = make_arg((g, 1, n)) + yield opinfo_core.SampleInput(self_t, args=(mat2_t, None, bias_t)) + + # Test with bias and out_dtype + if dtype in (torch.float16, torch.bfloat16): + yield opinfo_core.SampleInput( + self_t, args=(mat2_t, None, bias_t, torch.float32) + ) + def _mock_grouped_mm(self, mat2, offs=None, bias=None, out_dtype=None): + if hasattr(torch.ops.aten, "_grouped_mm"): + try: + return torch.ops.aten._grouped_mm(self, mat2, offs, bias, out_dtype) + except Exception: + pass res = torch.matmul(self, mat2) if bias is not None: res = res + bias + if out_dtype is not None: + res = res.to(out_dtype) return res