Select elements of numpy array via mask and preserving original dimensions

Question

Hello I have following data

ids = np.concatenate([1.0 * np.ones(shape=(4, 9,)), 
                      2.0 * np.ones(shape=(4, 3,))], axis=1)

logits = np.ran