What is the broadcasting order here?

Given the following:

polygons = np.array([
       [[1, 0],
        [1, 1],
        [0, 0],
        [0, 1]],

       [[3, 2],
        [2, 2],
        [4, 4],
        [2, 3]]])

sort_idx = np.array([
       [2, 0, 1, 3],
       [1, 0, 2, 3]])

polygons[sort_idx]  #< The problematic part

The expected output is:

array([[[0, 0],
        [1, 0],
        [1, 1],
        [0, 1]],

       [[2, 2],
        [3, 2],
        [4, 4],
        [2, 3]]])

I expect polygons[sort_idx] to return the "sorted" rows of polygons, with the sorting order given by the values in sort_idx.
polygons[0][sorted_idx[0]] (and the equivalent for [1]) give the correct values, but running polygons[sort_idx] raises:

IndexError: index 2 is out of bounds for axis 0 with size 2

I estimate the issue is connected to me not understanding something of how the operation is broadcasted, but I don’t really know what to search for or how to phrase the question. I saw similar questions recommending to use np.take() or polygons[sort_idx,...], but both raise the same error. What am I missing?

>Solution :

Use np.take_along_axis:

np.take_along_axis(polygons, sort_idx[..., None], 1)
Out[]: 
array([[[0, 0],
        [1, 0],
        [1, 1],
        [0, 1]],

       [[2, 2],
        [3, 2],
        [4, 4],
        [2, 3]]])