Take_along_dim vs gather

Hi I am wondering what’s the difference between the newly added take_along_dim & gather? It seems that they are pretty much similar other than take_along_dim can have user not specifying dim parameter. In this case seems that torch can find the best broadcast, for example I tried logits of shape (N, K) and use index of logits.argmax(-1) i.e. shape (N,) it can still find the correct broadcast (N, 1), which I believe is not possible by the traditional numpy broadcasting.


This method was adden in this PR as part of the numpy compatibility workload and thus might be similar to gather.

1 Like