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 577e57641d [Relax] Fix bucketize output dtype during legalization
(#19936)
577e57641d is described below
commit 577e57641d46db11ea3fce348e83e16417578748
Author: Hangshuai He <[email protected]>
AuthorDate: Fri Jul 10 12:33:28 2026 +0800
[Relax] Fix bucketize output dtype during legalization (#19936)
This PR fixes Relax bucketize lowering to pass the correct integer
output dtype to TOPI searchsorted.
Previously, both LegalizeOps and DispatchSortScan passed the input
tensor dtype as the output dtype. For float input tensors, this caused
TOPI searchsorted to receive a float output dtype, which later failed
during binary-search lowering because bucket indices must be integer
values.
This patch derives the output dtype from bucketize's out_int32
attribute:
- int32 when out_int32=True
- int64 otherwise
A numerical ExportedProgram frontend test is added to cover:
- right=False
- right=True
- out_int32=False
- out_int32=True
- float input values on bucket boundaries
Test:
python -m pytest
tests/python/relax/test_frontend_from_exported_program.py -k "bucketize"
-q
---
python/tvm/relax/backend/dispatch_sort_scan.py | 3 ++-
python/tvm/relax/transform/legalize_ops/search.py | 3 ++-
tests/python/relax/test_frontend_from_exported_program.py | 13 +++++++++++++
3 files changed, 17 insertions(+), 2 deletions(-)
diff --git a/python/tvm/relax/backend/dispatch_sort_scan.py
b/python/tvm/relax/backend/dispatch_sort_scan.py
index 20550870ce..929e2e2781 100644
--- a/python/tvm/relax/backend/dispatch_sort_scan.py
+++ b/python/tvm/relax/backend/dispatch_sort_scan.py
@@ -85,8 +85,9 @@ class SortScanDispatcher(BackendDispatcher):
with tgt:
if self.is_gpu_target(tgt):
te_func = topi.gpu.searchsorted
+ out_dtype = "int32" if call.attrs.out_int32 else "int64"
return self.builder_.call_te(
- te_func, boundaries, input_tensor, right,
input_tensor.ty.dtype.dtype
+ te_func, boundaries, input_tensor, right, out_dtype
)
if call.op.name == "relax.sort":
tgt = self._get_target(call.ty)
diff --git a/python/tvm/relax/transform/legalize_ops/search.py
b/python/tvm/relax/transform/legalize_ops/search.py
index 36e04d5c59..fdbbae47fc 100644
--- a/python/tvm/relax/transform/legalize_ops/search.py
+++ b/python/tvm/relax/transform/legalize_ops/search.py
@@ -48,4 +48,5 @@ def _bucketize(bb, call):
input_tensor = call.args[0]
boundaries = call.args[1]
right = call.attrs.right
- return bb.call_te(topi.searchsorted, boundaries, input_tensor, right,
input_tensor.ty.dtype)
+ out_dtype = "int32" if call.attrs.out_int32 else "int64"
+ return bb.call_te(topi.searchsorted, boundaries, input_tensor, right,
out_dtype)
diff --git a/tests/python/relax/test_frontend_from_exported_program.py
b/tests/python/relax/test_frontend_from_exported_program.py
index e4200206a6..f093cf324d 100644
--- a/tests/python/relax/test_frontend_from_exported_program.py
+++ b/tests/python/relax/test_frontend_from_exported_program.py
@@ -7845,6 +7845,19 @@ def test_bucketize():
verify_model(Bucketize(), (input_tensor, boundaries), {}, Expected)
[email protected]("right", [False, True])
[email protected]("out_int32", [False, True])
+def test_bucketize_numerically(right, out_int32):
+ class Bucketize(Module):
+ def forward(self, input_tensor, boundaries):
+ return torch.bucketize(input_tensor, boundaries, right=right,
out_int32=out_int32)
+
+ input_tensor = torch.tensor([-0.5, 0.0, 0.5, 1.0, 2.0, 2.5],
dtype=torch.float32)
+ boundaries = torch.tensor([0.0, 1.0, 2.0], dtype=torch.float32)
+
+ verify_model_numerically(Bucketize(), (input_tensor, boundaries))
+
+
def test_argsort():
class Argsort(Module):
def forward(self, x):
