[routines] Re-write argmax and argmin to returns indices along the given axis
#230
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
This pull request focuses on enhancing the functionality of the
NDArrayclass and improving theargmaxandargminmethods. Additionally, new tests have been added to ensure the correctness of these methods.Improvements to
argmaxandargminmethods:numojo/routines/searching.mojo: Introduced newargmaxandargminfunctions that handle both 1D and multi-dimensional arrays. These functions now support axis-based operations and include detailed documentation and examples.numojo/routines/functional.mojo: Added overloads ofapply_along_axisto support the newargmaxandargminfunctions, ensuring they reduce the dimension of the input array appropriately.Enhancements to
NDArrayclass:numojo/core/ndarray.mojo: Theargmaxandargminmethods have been updated to use the newsearching.argmaxandsearching.argminfunctions. These methods now return indices of the maximum and minimum values along an axis, with improved documentation.Bug fixes:
numojo/routines/math/extrema.mojo: Corrected error messages in themaxandminfunctions to accurately reflect the function names.New tests:
tests/routines/test_searching.mojo: Added comprehensive tests for theargmaxandargminfunctions, including tests for 1D, 2D, and 3D arrays, as well as tests for negative axes and Fortran-order arrays. Also included tests fortake_along_axiswithargmaxandargmin.