[Python] How to broadcast np.take_along_axis?

My objective is to sort Y first (increasing, from 0 to 3), then sort the X according to the sorted Y. Note that X and Y are one to one.

Y = np.random.randint(0, 4, (2, 10,))
print("before Y", Y.shape) # before Y (2, 10)
print(Y)
print()

idx = np.argsort(Y, axis=1)

Y_ordered = np.take_along_axis(Y, idx, axis=1)
print("reordered Y", Y_ordered.shape) # reordered Y (2, 10)
print(Y_ordered)
print()

###

X = np.random.rand(2, 10, 64, 64)

X_ordered = np.take_along_axis(X, idx, axis=1) # ValueError: `indices` and `arr` must have the same number of dimensions
print("reordered X", X_ordered.shape)
print(X_ordered)
print()


I got the Y_ordered that I want, but how to broadcast the same idx on another array X with different shape, but have the same size in the first 2 dimensions.

Continue reading...