This is an automated email from the ASF dual-hosted git repository.
tlopex pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new ca6fcc25b5 [BugFix][Relax] Preserve take mode in
ReorderTakeAfterMatmul (#20206)
ca6fcc25b5 is described below
commit ca6fcc25b5655a92c25794515c9b13447ac9b91b
Author: Midst <[email protected]>
AuthorDate: Fri Aug 28 15:05:29 2026 +0800
[BugFix][Relax] Preserve take mode in ReorderTakeAfterMatmul (#20206)
`ReorderTakeAfterMatmul` rebuilds `relax.take` calls without forwarding
the original `mode`, so `mode="clip"` silently becomes the default
`mode="fast"`. This changes the defined behavior for out-of-bounds
indices.
This change forwards `TakeAttrs::mode` in both the simple and batched
weight rewrite paths, and adds structural regression tests for both
paths.
Testing:
- `git diff --check`
- C++ syntax checks with TVM-generated compiler flags
- `python -m py_compile
tests/python/relax/test_transform_reorder_take_after_matmul.py`
Full pytest was not run because the current Windows host does not have
MSVC. The available MinGW build stops in TVM-FFI's Windows platform code
before compiling this transform.
Fixes #20201
---
src/relax/transform/reorder_take_after_matmul.cc | 4 +-
.../test_transform_reorder_take_after_matmul.py | 66 ++++++++++++++++++++++
2 files changed, 68 insertions(+), 2 deletions(-)
diff --git a/src/relax/transform/reorder_take_after_matmul.cc
b/src/relax/transform/reorder_take_after_matmul.cc
index f693d0278b..19e30f8c68 100644
--- a/src/relax/transform/reorder_take_after_matmul.cc
+++ b/src/relax/transform/reorder_take_after_matmul.cc
@@ -94,7 +94,7 @@ std::tuple<DFPattern, ffi::TypedFunction<Expr(Expr,
ffi::Map<DFPattern, Expr>)>>
// out_table.shape = [*batch, table_size]
auto out_table = matmul(lhs, weights, std::nullopt);
// new_output.shape = [*batch, outfeatures]
- auto new_output = take(out_table, indices, matmul_ty->ndim - 1);
+ auto new_output = take(out_table, indices, matmul_ty->ndim - 1,
attrs->mode);
return new_output;
} else if (lhs_ty->ndim == 3 && weights_ty->ndim == 3 && indices_ty->ndim
== 1 && axis == 0 &&
@@ -130,7 +130,7 @@ std::tuple<DFPattern, ffi::TypedFunction<Expr(Expr,
ffi::Map<DFPattern, Expr>)>>
// operations.
// duplicated_output.shape = [batch1, batch2, batch1, outfeatures]
- auto duplicated_output = take(indexed_output, indices, 2);
+ auto duplicated_output = take(indexed_output, indices, 2, attrs->mode);
// new_output.shape = [batch1, batch2, outfeatures]
auto new_output = einsum(Tuple({duplicated_output}), "ijik->ijk");
diff --git a/tests/python/relax/test_transform_reorder_take_after_matmul.py
b/tests/python/relax/test_transform_reorder_take_after_matmul.py
index 8e7243d02e..3d2de09907 100644
--- a/tests/python/relax/test_transform_reorder_take_after_matmul.py
+++ b/tests/python/relax/test_transform_reorder_take_after_matmul.py
@@ -185,5 +185,71 @@ class TestDynamicBatchedActivationsAndWeights(Base):
return out
+class TestPreserveTakeMode(Base):
+ @I.ir_module
+ class Before:
+ @R.function
+ def main(
+ x: R.Tensor([1, 16], "float32"),
+ weight_table: R.Tensor([16, 64], "float32"),
+ routing_table: R.Tensor([32], "int64"),
+ ) -> R.Tensor([1, 32], "float32"):
+ with R.dataflow():
+ weight = R.take(weight_table, routing_table, axis=1,
mode="clip")
+ out = R.matmul(x, weight)
+ R.output(out)
+ return out
+
+ @I.ir_module
+ class Expected:
+ @R.function
+ def main(
+ x: R.Tensor([1, 16], "float32"),
+ weight_table: R.Tensor([16, 64], "float32"),
+ routing_table: R.Tensor([32], "int64"),
+ ) -> R.Tensor([1, 32], "float32"):
+ with R.dataflow():
+ out_table = R.matmul(x, weight_table)
+ out = R.take(out_table, routing_table, axis=1, mode="clip")
+ R.output(out)
+ return out
+
+
+class TestPreserveTakeModeForBatchedWeights(Base):
+ @I.ir_module
+ class Before:
+ @R.function
+ def main(
+ x: R.Tensor([128, 1, 16], "float32"),
+ weight_table: R.Tensor([64, 16, 32], "float32"),
+ routing_table: R.Tensor([128], "int64"),
+ ) -> R.Tensor([128, 1, 32], "float32"):
+ with R.dataflow():
+ weight = R.take(weight_table, routing_table, axis=0,
mode="clip")
+ out = R.matmul(x, weight)
+ R.output(out)
+ return out
+
+ @I.ir_module
+ class Expected:
+ @R.function
+ def main(
+ x: R.Tensor([128, 1, 16], "float32"),
+ weight_table: R.Tensor([64, 16, 32], "float32"),
+ routing_table: R.Tensor([128], "int64"),
+ ) -> R.Tensor([128, 1, 32], "float32"):
+ with R.dataflow():
+ reordered_weight = R.permute_dims(weight_table, [1, 0, 2])
+ fused_weight = R.reshape(reordered_weight, [16, 2048])
+ fused_output = R.matmul(x, fused_weight)
+ reordered_output = R.reshape(fused_output, [128, 1, 64, 32])
+ tabular_output = R.take(
+ reordered_output, routing_table, axis=2, mode="clip"
+ )
+ out = R.einsum([tabular_output], "ijik->ijk")
+ R.output(out)
+ return out
+
+
if __name__ == "__main__":
tvm.testing.main()
