Port `amax` to structured kernel (#72124)
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/72124
Reviewed By: bdhirsh
Differential Revision: D34215708
Pulled By: ansley
fbshipit-source-id: fee887e331cb8bd9fab3d9d958ff13ac8d07be27
(cherry picked from commit 94dbb5b7e7e14a663dc02ecf5013fad10b8701b3)
