[resubmit] Add amax/amin (#43819)
Summary:
Resubmit for landing next week.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/43819
Reviewed By: ngimel
Differential Revision: D23421906
Pulled By: mruberry
fbshipit-source-id: 23dd60d1e365bb1197d660c3bfad7ee07ba3e97f