Skip to content

Conversation

@forfudan
Copy link
Collaborator

This pull request focuses on enhancing the functionality of the NDArray class and improving the argmax and argmin methods. Additionally, new tests have been added to ensure the correctness of these methods.

Improvements to argmax and argmin methods:

  • numojo/routines/searching.mojo: Introduced new argmax and argmin functions 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 of apply_along_axis to support the new argmax and argmin functions, ensuring they reduce the dimension of the input array appropriately.

Enhancements to NDArray class:

  • numojo/core/ndarray.mojo: The argmax and argmin methods have been updated to use the new searching.argmax and searching.argmin functions. 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 the max and min functions to accurately reflect the function names.

New tests:

  • tests/routines/test_searching.mojo: Added comprehensive tests for the argmax and argmin functions, including tests for 1D, 2D, and 3D arrays, as well as tests for negative axes and Fortran-order arrays. Also included tests for take_along_axis with argmax and argmin.

@forfudan forfudan requested a review from shivasankarka March 21, 2025 20:22
@MadAlex1997 MadAlex1997 merged commit 4815859 into Mojo-Numerics-and-Algorithms-group:pre-0.7 Mar 21, 2025
2 checks passed
@forfudan forfudan deleted the arg branch March 21, 2025 22:02
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants