Update argmax.py

[timmens]

In this PR, we update everything related to argmax.py.

• [x] Add type hints (See JAX Typing Best Practices)
• [x] Run mypy and add it to CI
• [x] Add an explanation of what argmax and segment_argmax do to the developer notes
  • [x] Briefly explain what the functions do on a programming level
  • [x] Draw a connection between the math notation of the dynamic programming model and functions

[codecov[bot]]

Codecov Report

All modified and coverable lines are covered by tests :white_check_mark:

Project coverage is 98.31%. Comparing base (56f244d) to head (7600ad3). Report is 1 commits behind head on main.

:exclamation: Current head 7600ad3 differs from pull request most recent head 3dcc8a2. Consider uploading reports for the commit 3dcc8a2 to get more accurate results

Additional details and impacted files

```diff @@ Coverage Diff @@ ## main #62 +/- ## ======================================= Coverage 98.30% 98.31% ======================================= Files 42 42 Lines 2415 2429 +14 ======================================= + Hits 2374 2388 +14 Misses 41 41 ```

:umbrella: View full report in Codecov by Sentry.
:loudspeaker: Have feedback on the report? Share it here.