Thank you very much for your help. After a lot of Googling and also discussions with ChatGPT, my understanding is that one has to look at this organization of elements, rather like a coordinate system with N dimensions (unlike Cartesian system which has only 2 dimensions) where a point (= an element in the array) is defined like (x_0, x_1, x_2, … x_i, …, x_N-1) for all i in {1, 2, …, N-1}. In that case, sorting along a specific axis, means to increment indices along that specific axis, in a way that in each iteration, elements in corresponding positions (= elements that would be compared with each other during the sort operation):
-
Have exactly the same value on each pair of non-sorting axis (of course, that doesn’t mean that they have necessarily the same value, on all not sorting axis. All that matters, is to have the same value on each pair of non sorting axis, when we compare each pair of them)
-
And they differ only on their index value of the sorting axis.
By this logic, it makes perfectly sense and if we read indices of the argsort, column by column from the top to the bottom, each of those indices are indeed the index of the smallest (according to the sort order) element (again, very important to note: smallest only in the context of corresponding positions) picked during the sort operation.
Just as an example, looking to the array elements in my above-mentioned post, using this method for analyzing elements, assuming that we sort along axis=1, we can say that -28 and 69 are in corresponding positions (= they will be compared with each other during the sort operation)
>>> a[0, :, 0, 0]
array([-28, 69])
>>>
>>>
>>> # We can also look at them separately
>>> a[0, 0, 0, 0]
np.int64(-28)
>>>
>>>
>>> a[0, 1, 0, 0]
np.int64(69)
Now, how exactly take_along_axis figures out intelligently the indices according to the initial shape, I’ve no idea! I tried to look a bit on Github into the source code of NumPy, but it seems to me that it is far to be just a simple small function and a lot of imported modules are implicated. So far, what I have understood is that unless the sort operation and the original array are 1-D arrays, one has to use take_along_axis to obtain correctly the initial structure. Thaks again for your help.