ENH : np.argmax is unusually time consuming for multidimensional array

Describe the issue:

Both np.amax and np.argmax are expected to have a time complexity of O(n)，so they should have similar computational times.
However, they only exhibit comparable performance on 1D arrays.
For 2D or higher dimensional arrays, np.amax consistently outperforms np.argmax by a factor of 8x or more. It's strange.

Reproduce the code example:

import timeit

stmt1 = "np.argmax(a, axis=0)"
stmt2 = "np.amax(a, axis=0)"
setup_1d = "import numpy as np; a = np.random.rand(3*768*768)"
setup_2d = "import numpy as np; a = np.random.rand(3,768*768)"

execution_time1 = timeit.timeit(stmt1, setup=setup_2d, number=1000)
print(f"Execution time for np.argmax on 2d array: {execution_time1} seconds")

execution_time2 = timeit.timeit(stmt2, setup=setup_2d, number=1000)
print(f"Execution time for np.amax on 2d array: {execution_time2} seconds")

execution_time1 = timeit.timeit(stmt1, setup=setup_1d, number=1000)
print(f"Execution time for np.argmax on 1d array: {execution_time1} seconds")

execution_time2 = timeit.timeit(stmt2, setup=setup_1d, number=1000)
print(f"Execution time for np.amax on 1d array: {execution_time2} seconds")

Error message:

Execution time for np.argmax on 2d array: 16.13085489999503 seconds
Execution time for np.amax on 2d array: 2.400201399810612 seconds
Execution time for np.argmax on 1d array: 0.6763406000100076 seconds
Execution time for np.amax on 1d array: 0.4886799002997577 seconds

Python and NumPy Versions:

Python: 3.10.13
Numpy: 1.26.4

Runtime Environment:

No response

Context for the issue:

No response

Start by running the provided timeit reproduction with NumPy 1.26.4, comparing np.argmax and np.amax on the one- and two-dimensional arrays. Trace the np.argmax and np.amax entry points to identify where multidimensional execution diverges. Done means the reported performance gap is explained and, if actionable, reduced without changing results.