Reduce rank and set values according to 3rd dimension

@SidneyLann is this for weighted loss? If not why not make that a placeholder?