[ROCm] Enable sort operator BF16 support (#72854)
Summary:
The changes add support for dtype BF16 for sort operator in ROCm.
Relates - https://github.com/pytorch/pytorch/pull/58196
Relanding the change - https://github.com/pytorch/pytorch/pull/71226
jeffdaily jithunnair-amd dllehr-amd Please review this PR.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/72854
Reviewed By: zou3519
Differential Revision: D34284313
Pulled By: malfet
fbshipit-source-id: abcfea84ea53874008d56416425849e990ebf15b
(cherry picked from commit e9e7e3e0472b726ff2fd5f115962d3c835fb33db)