Add amax

Details:
- Add amax and tests
- Fix minor typos