Port `argmax` to structured kernels. (#59937)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/59937
Tracking issue: #55070
Test Plan: Imported from OSS
Reviewed By: soulitzer
Differential Revision: D29104397
Pulled By: ezyang
fbshipit-source-id: 580355cf3b4e9e5c934b4e51a16196087bcb3459
