[dist_optim] Add distributed functional Adamax optimizer (#55833)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/55833
Add distributed functional Adamax optimizer, to support in TorchScript
ghstack-source-id: 126325538
Reviewed By: rohan-varma
Differential Revision: D26696540
fbshipit-source-id: 6242faebd2476847831a05df7f8b0d616f2b5355
