Port `argmax` to structured kernels. (#60363)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/60363
Tracking issue: #55070
This PR was openned so as to solve the CI failures in main when merging: #59371 #59372 #59373 #59937 #59938.
Test Plan: Imported from OSS
Reviewed By: albanD
Differential Revision: D29265857
Pulled By: ezyang
fbshipit-source-id: 586914d2aa79028c56988896093945755a2b9781
