row-wise indexing in numpy,"A[np.arange(2)[:, (None)], B]"