[Relax] Fix bucketize output dtype during legalization

[V-aerus]

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

[gemini-code-assist]

Code Review

This pull request updates the legalization and dispatching of the bucketize and searchsorted operations in TVM Relax to correctly handle the out_int32 attribute, determining whether the output data type should be int32 or int64 instead of using the input tensor's data type. Additionally, a new parameterized test has been added to verify the numerical correctness of torch.bucketize with various configurations of right and out_int32. I have no feedback to provide as the changes are correct and well-tested.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.