In Numpy, how can I index an B x N x M array with a size B array of indices to yield a B x M array?
Figured it out. Say the BxNxM array is called A, then the desired output is A[(np.arange(B), indices)].
I want to index the 2nd axis of a BxNxM array (the N axis) with B indices and get a BxM array as the result.