import torch


def op_norm(outputs, op_threshs):
    """normalize outputs according to operating points for a given model.
    Args: 
        outputs: outputs of self.classifier(). torch.Size(batch_size, num_tasks) 
        op_threshs_arr: torch.Size(batch_size, num_tasks) with self.op_threshs expanded.
    Returns:
        outputs_new: normalized outputs, torch.Size(batch_size, num_tasks)
    """
    op_threshs = op_threshs.expand(outputs.shape[0],-1)
    outputs_new = torch.zeros(outputs.shape, device = outputs.device)
    mask_leq = outputs<op_threshs
    outputs_new[mask_leq] = outputs[mask_leq]/(op_threshs[mask_leq]*2)
    outputs_new[~mask_leq] = 1.0 - ((1.0 - outputs[~mask_leq])/((1-op_threshs[~mask_leq])*2))
    
    return outputs_new