From fa0f1ad9fe7ced1a0c6c7c7fef6017df2d72e900 Mon Sep 17 00:00:00 2001
From: Hehang Shuai <152275085+V-aerus@users.noreply.github.com>
Date: Sat, 4 Jul 2026 14:39:03 +0800
Subject: [PATCH] [Relax] Fix bucketize output dtype during legalization

---
 python/tvm/relax/backend/dispatch_sort_scan.py      |  3 ++-
 python/tvm/relax/transform/legalize_ops/search.py   |  3 ++-
 .../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 20550870cef0..929e2e278138 100644
--- a/python/tvm/relax/backend/dispatch_sort_scan.py
+++ b/python/tvm/relax/backend/dispatch_sort_scan.py
@@ -85,8 +85,9 @@ def visit_call_(self, call: relax.Call) -> relax.Expr:
             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 36e04d5c59d4..fdbbae47fc8c 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 d91aa46f219b..c122a7d51679 100644
--- a/tests/python/relax/test_frontend_from_exported_program.py
+++ b/tests/python/relax/test_frontend_from_exported_program.py
@@ -7829,6 +7829,19 @@ def main(
     verify_model(Bucketize(), (input_tensor, boundaries), {}, Expected)
 
 
+@pytest.mark.parametrize("right", [False, True])
+@pytest.mark.parametrize("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):