whoops forgot to set partition_step in AddReduceSingle, to be passed to as_mask