Numpy: select along axis of 3d array using 2d array

Question

I\'ve been struggling with this for a few hours and can\'t quite get my head around it. The setup is something like this:

    A.shape # (T,N,K)
    B.shape #