Skip to content

[Fix][Relax] Preserve out_dtype in matmul rewrite passes - #20331

Open
tintin1942 wants to merge 1 commit into
apache:mainfrom
tintin1942:fix/relax-preserve-matmul-out-dtype
Open

tintin1942 wants to merge 1 commit into
apache:mainfrom
tintin1942:fix/relax-preserve-matmul-out-dtype

Conversation

@tintin1942

Copy link
Copy Markdown

This PR fixes ExpandMatmulOfSum and ReorderTakeAfterMatmul dropping an explicitly specified out_dtype when constructing replacement matmul calls. With float16 operands and out_dtype="float32", the transformed functions incorrectly returned float16.

The rewrite passes constructed the replacement matmul calls with an unspecified output dtype. This change forwards the original MatmulAttrs::out_dtype in:

  • Both replacement matmuls created by ExpandMatmulOfSum.
  • The simple and batched-weight paths of ReorderTakeAfterMatmul.

Structural regression tests cover all three rewrite paths with float16 inputs and an explicit float32 output dtype.

Tests:

  • cmake --build build --parallel 12
  • python -m pytest -p no:rerunfailures tests/python/relax/test_transform_expand_matmul_of_sum.py tests/python/relax/test_transform_reorder_take_after_matmul.py tests/python/relax/test_op_linear_algebra.py -q — 26 passed
  • ctest --test-dir build --output-on-failure — 137/137 passed
  • pre-commit run --files src/relax/transform/expand_matmul_of_sum.cc src/relax/transform/reorder_take_after_matmul.cc tests/python/relax/test_transform_expand_matmul_of_sum.py tests/python/relax/test_transform_reorder_take_after_matmul.py — passed

Fixes #20198

Forward the original MatmulAttrs::out_dtype when ExpandMatmulOfSum and ReorderTakeAfterMatmul construct replacement matmul calls.

Add regression coverage for the expansion and both take-reordering paths.
@tintin1942 tintin1942 changed the title [BugFix][Relax] Preserve out_dtype in matmul rewrite passes [Fix][Relax] Preserve out_dtype in matmul rewrite passes Sep 13, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Bug][Relax] Matmul rewrite passes drop explicit out_dtype and change float32 outputs to float16

1 participant