[torchlib] Support offset-based grouped_mm layouts

[TANGBUDU]

Summary

Builds on the aten_grouped_mm converter added in #2805 and adds support for its offsets form. The existing dense, bias, and cast paths are unchanged.

The converter now handles three layouts:

• 2D x 3D: split rows of the first operand using the offsets.
• 3D x 2D: split columns of the second operand using the offsets.
• 2D x 2D: split the contraction dimension and stack the group results.

The number of groups must be known at export time, but offset values and input lengths can remain dynamic. If the final offset does not cover the full input, the converter keeps PyTorch's full output shape and fills the unwritten tail with zeros.

The offsets path does not currently support bias or changing the output dtype. The existing dense converter still supports both.

Testing

OMP_NUM_THREADS=1 OPENBLAS_NUM_THREADS=1 MKL_NUM_THREADS=1 \
python -m pytest \
  tests/function_libs/torch_lib/grouped_mm_test.py \
  tests/function_libs/torch_lib/ops_test.py \
  -k grouped_mm -q

• PyTorch 2.14 environment: 50 passed, 1 skipped.
• Repository-pinned environment: 46 passed, 5 skipped.
• Tests cover all three layouts, uneven and empty groups, dynamic offsets, dynamic input lengths, and a single group.
• Three dynamic BF16 exports pass the ONNX checker. BF16 runtime and GPU execution were not tested.
• lintrunner -a and git diff --check pass.

AI assistance

ChatGPT assisted with the implementation, tests, and result analysis. I reviewed the changes and the test results above.

[TANGBUDU]

@microsoft-github-policy-service agree

[codecov]

Codecov Report

❌ Patch coverage is 95.23810% with 2 lines in your changes missing coverage. Please review.
✅ Project coverage is 72.76%. Comparing base (d1c005d) to head (8e000b2).
⚠️ Report is 1 commits behind head on main.

Files with missing lines	Patch %	Lines
onnxscript/function_libs/torch_lib/ops/core.py	95.23%	1 Missing and 1 partial ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #3017      +/-   ##
==========================================
+ Coverage   72.70%   72.76%   +0.05%     
==========================================
  Files         265      265              
  Lines       32298    32339      +41     
  Branches     3059     3072      +13     
==========================================
+ Hits        23481    23530      +49     
+ Misses       7779     7770       -9     
- Partials     1038     1039       +1     

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

[TANGBUDU]

Justin Chu (@justinchuby) Hi, now that #2805 is merged, this PR adds the offset-based layouts on top of it. It also fixes the _grouped_mm output-consistency test that currently fails on main: the old samples use strides that PyTorch's _grouped_mm rejects ("strides should be multiple of 16 bytes"), so they are replaced with aligned shapes. The bias and out_dtype samples were dropped because PyTorch's CPU _grouped_mm doesn't support them ("Bias not supported yet"); that path is still tested in grouped_mm_test.py against NumPy.

I merged the latest main today, and the CI runs are waiting for approval. Could you review it when you have a chance? Thanks!