The inner loop that would update the broadcastable pattern according
to the rebroadcast dict would overwrite the idx extracted from the
Subtensor, thus always selecting the last dim.