jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def get_t(N, Ca, C, eps=1e-5):
    #N, Ca, C - [I, B, L, 3]
    #R - [I, B, L, 3, 3], det(R)=1, inv(R) = R.T, R is a rotation matrix
    #t - [I, B, L, L, 3] is the global rotation and translation invariant displacement
    v1 = N-Ca # (I, B, L, 3)
    v2 = C-Ca # (I, B, L, 3)
    e1 = v1/(torch.norm(v1, dim=-1, keepdim=True)+eps) # (I, B, L, 3)
    u2 = v2 - (torch.einsum('iblj,iblj->ibl', v2, e1).unsqueeze(-1)*e1) # (I,B,L,3)
    e2 = u2/(torch.norm(u2, dim=-1, keepdim=True)+eps) # (I, B, L, 3)
    e3 = torch.cross(e1, e2, dim=-1) # (I, B, L, 3)
    R = torch.stack((e1,e2,e3), dim=-2) #[I,B,L,3,3] - rotation matrix
    t = Ca.unsqueeze(-2) - Ca.unsqueeze(-3) # (I,B,L,L,3)
    t = torch.einsum('ibljk, iblmk -> iblmj', R, t) # (I,B,L,L,3)
    return t
//  [python]
def get_t(N, Ca, C, eps=1e-5):
    #N, Ca, C - [I, B, L, 3]
    #R - [I, B, L, 3, 3], det(R)=1, inv(R) = R.T, R is a rotation matrix
    #t - [I, B, L, L, 3] is the global rotation and translation invariant displacement
    v1 = N-Ca # (I, B, L, 3)
    v2 = C-Ca # (I, B, L, 3)
    e1 = v1/(torch.linalg.norm(v1, dim=-1, keepdim=True)+eps) # (I, B, L, 3)
    u2 = v2 - (torch.einsum('iblj,iblj->ibl', v2, e1).unsqueeze(-1)*e1) # (I,B,L,3)
    e2 = u2/(torch.linalg.norm(u2, dim=-1, keepdim=True)+eps) # (I, B, L, 3)
    e3 = torch.cross(e1, e2, dim=-1) # (I, B, L, 3)
    R = torch.stack((e1,e2,e3), dim=-2) #[I,B,L,3,3] - rotation matrix
    t = Ca.unsqueeze(-2) - Ca.unsqueeze(-3) # (I,B,L,L,3)
    t = torch.einsum('ibljk, iblmk -> iblmj', R, t) # (I,B,L,L,3)
    return t
//  [python]
def test_jacobian_pdiag_vs_pdense():
    for get_task in nonlinear_tasks:
        loader, lc, parameters, model, function, n_output = get_task()
        generator = Jacobian(
            layer_collection=lc, model=model, function=function, n_output=n_output
        )
        PMat_diag = PMatDiag(generator=generator, examples=loader)
        PMat_dense = PMatDense(generator=generator, examples=loader)
        dw = random_pvector(lc, device=device)

        # Test get_dense_tensor
        matrix_diag = PMat_diag.get_dense_tensor()
        matrix_dense = PMat_dense.get_dense_tensor()
        check_tensors(torch.diag(matrix_diag), torch.diag(matrix_dense))
        assert torch.norm(matrix_diag - torch.diag(torch.diag(matrix_diag))) < 1e-5

        # Test trace
        check_ratio(torch.trace(matrix_diag), PMat_diag.trace())

        # Test frobenius
        check_ratio(torch.norm(matrix_diag), PMat_diag.frobenius_norm())

        # Test mv
        mv_direct = torch.mv(matrix_diag, dw.get_flat_representation())
        mv_PMat_diag = PMat_diag.mv(dw)
        check_tensors(mv_direct, mv_PMat_diag.get_flat_representation())

        # Test vTMv
        vTMv_direct = torch.dot(mv_direct, dw.get_flat_representation())
        vTMv_PMat_diag = PMat_diag.vTMv(dw)
        check_ratio(vTMv_direct, vTMv_PMat_diag)

        # Test inverse
        regul = 1e-3
        PMat_diag_inverse = PMat_diag.inverse(regul)
        prod = torch.mm(
            matrix_diag + regul * torch.eye(lc.numel(), device=device),
            PMat_diag_inverse.get_dense_tensor(),
        )
        check_tensors(torch.eye(lc.numel(), device=device), prod)

        # Test solve
        regul = 1e-3
        Mv_regul = torch.mv(
            matrix_diag + regul * torch.eye(PMat_diag.size(0), device=device),
            dw.get_flat_representation(),
        )
        Mv_regul = PVector(layer_collection=lc, vector_repr=Mv_regul)
        dw_using_inv = PMat_diag.solve(Mv_regul, regul=regul)
        check_tensors(
            dw.get_flat_representation(),
            dw_using_inv.get_flat_representation(),
            eps=5e-3,
        )

        # Test get_diag
        diag_direct = torch.diag(matrix_diag)
        diag_PMat_diag = PMat_diag.get_diag()
        check_tensors(diag_direct, diag_PMat_diag)

        # Test add, sub, rmul
        loader, lc, parameters, model, function, n_output = get_task()
        generator = Jacobian(
            layer_collection=lc, model=model, function=function, n_output=n_output
        )
        PMat_diag2 = PMatDiag(generator=generator, examples=loader)

        check_tensors(
            PMat_diag.get_dense_tensor() + PMat_diag2.get_dense_tensor(),
            (PMat_diag + PMat_diag2).get_dense_tensor(),
        )
        check_tensors(
            PMat_diag.get_dense_tensor() - PMat_diag2.get_dense_tensor(),
            (PMat_diag - PMat_diag2).get_dense_tensor(),
        )
        check_tensors(
            1.23 * PMat_diag.get_dense_tensor(), (1.23 * PMat_diag).get_dense_tensor()
        )
//  [python]
def test_jacobian_pdiag_vs_pdense():
    for get_task in nonlinear_tasks:
        loader, lc, parameters, model, function, n_output = get_task()
        generator = Jacobian(
            layer_collection=lc, model=model, function=function, n_output=n_output
        )
        PMat_diag = PMatDiag(generator=generator, examples=loader)
        PMat_dense = PMatDense(generator=generator, examples=loader)
        dw = random_pvector(lc, device=device)

        # Test get_dense_tensor
        matrix_diag = PMat_diag.get_dense_tensor()
        matrix_dense = PMat_dense.get_dense_tensor()
        check_tensors(torch.diag(matrix_diag), torch.diag(matrix_dense))
        assert torch.linalg.norm(matrix_diag - torch.diag(torch.diag(matrix_diag))) < 1e-5

        # Test trace
        check_ratio(torch.trace(matrix_diag), PMat_diag.trace())

        # Test frobenius
        check_ratio(torch.linalg.norm(matrix_diag), PMat_diag.frobenius_norm())

        # Test mv
        mv_direct = torch.mv(matrix_diag, dw.get_flat_representation())
        mv_PMat_diag = PMat_diag.mv(dw)
        check_tensors(mv_direct, mv_PMat_diag.get_flat_representation())

        # Test vTMv
        vTMv_direct = torch.dot(mv_direct, dw.get_flat_representation())
        vTMv_PMat_diag = PMat_diag.vTMv(dw)
        check_ratio(vTMv_direct, vTMv_PMat_diag)

        # Test inverse
        regul = 1e-3
        PMat_diag_inverse = PMat_diag.inverse(regul)
        prod = torch.mm(
            matrix_diag + regul * torch.eye(lc.numel(), device=device),
            PMat_diag_inverse.get_dense_tensor(),
        )
        check_tensors(torch.eye(lc.numel(), device=device), prod)

        # Test solve
        regul = 1e-3
        Mv_regul = torch.mv(
            matrix_diag + regul * torch.eye(PMat_diag.size(0), device=device),
            dw.get_flat_representation(),
        )
        Mv_regul = PVector(layer_collection=lc, vector_repr=Mv_regul)
        dw_using_inv = PMat_diag.solve(Mv_regul, regul=regul)
        check_tensors(
            dw.get_flat_representation(),
            dw_using_inv.get_flat_representation(),
            eps=5e-3,
        )

        # Test get_diag
        diag_direct = torch.diag(matrix_diag)
        diag_PMat_diag = PMat_diag.get_diag()
        check_tensors(diag_direct, diag_PMat_diag)

        # Test add, sub, rmul
        loader, lc, parameters, model, function, n_output = get_task()
        generator = Jacobian(
            layer_collection=lc, model=model, function=function, n_output=n_output
        )
        PMat_diag2 = PMatDiag(generator=generator, examples=loader)

        check_tensors(
            PMat_diag.get_dense_tensor() + PMat_diag2.get_dense_tensor(),
            (PMat_diag + PMat_diag2).get_dense_tensor(),
        )
        check_tensors(
            PMat_diag.get_dense_tensor() - PMat_diag2.get_dense_tensor(),
            (PMat_diag - PMat_diag2).get_dense_tensor(),
        )
        check_tensors(
            1.23 * PMat_diag.get_dense_tensor(), (1.23 * PMat_diag).get_dense_tensor()
        )
//  [python]
def angle(a, b, c, eps=1e-6):
    '''
    Calculate cos/sin angle between ab and cb
    a,b,c have shape of (B, L, 3)
    '''
    B,L = a.shape[:2]

    u1 = a-b
    u2 = c-b

    u1_norm = torch.norm(u1, dim=-1, keepdim=True) + eps
    u2_norm = torch.norm(u2, dim=-1, keepdim=True) + eps

    # normalize u1 & u2 --> make unit vector
    u1 = u1 / u1_norm
    u2 = u2 / u2_norm
    u1 = u1.reshape(B*L, 3)
    u2 = u2.reshape(B*L, 3)

    # sin_theta = norm(a cross b)/(norm(a)*norm(b))
    # cos_theta = norm(a dot b) / (norm(a)*norm(b))
    sin_theta = torch.norm(torch.cross(u1, u2, dim=1), dim=1, keepdim=True).reshape(B, L, 1) # (B,L,1)
    cos_theta = torch.matmul(u1[:,None,:], u2[:,:,None]).reshape(B, L, 1)
    
    return torch.cat([cos_theta, sin_theta], axis=-1)
//  [python]
def angle(a, b, c, eps=1e-6):
    '''
    Calculate cos/sin angle between ab and cb
    a,b,c have shape of (B, L, 3)
    '''
    B,L = a.shape[:2]

    u1 = a-b
    u2 = c-b

    u1_norm = torch.linalg.norm(u1, dim=-1, keepdim=True) + eps
    u2_norm = torch.linalg.norm(u2, dim=-1, keepdim=True) + eps

    # normalize u1 & u2 --> make unit vector
    u1 = u1 / u1_norm
    u2 = u2 / u2_norm
    u1 = u1.reshape(B*L, 3)
    u2 = u2.reshape(B*L, 3)

    # sin_theta = norm(a cross b)/(norm(a)*norm(b))
    # cos_theta = norm(a dot b) / (norm(a)*norm(b))
    sin_theta = torch.linalg.norm(torch.cross(u1, u2, dim=1), dim=1, keepdim=True).reshape(B, L, 1) # (B,L,1)
    cos_theta = torch.matmul(u1[:,None,:], u2[:,:,None]).reshape(B, L, 1)
    
    return torch.cat([cos_theta, sin_theta], axis=-1)
//  [python]
def quaternion_loss(predictions, targets):
    """
    Computes the quaternion loss between predicted and target quaternions.
    Args:
        predictions (torch.Tensor): predicted quaternions of shape (batch_size, 4)
        targets (torch.Tensor): target quaternions of shape (batch_size, 4)
    Returns:
        quaternion_loss (torch.Tensor): quaternion loss value
    """
    # Normalize the predictions and targets
    predictions_norm = predictions / torch.norm(predictions, dim=1, keepdim=True)
    targets_norm = targets / torch.norm(targets, dim=1, keepdim=True)
    
    # Compute the dot product between the normalized predictions and targets
    dot_product = torch.sum(predictions_norm * targets_norm, dim=1)
    
    # Compute the angle between the predictions and targets
    angle = 2 * torch.acos(torch.clamp(dot_product, min=-1, max=1))
    
    # Compute the quaternion loss
    quaternion_loss = torch.mean(angle)
    
    return quaternion_loss
//  [python]
def quaternion_loss(predictions, targets):
    """
    Computes the quaternion loss between predicted and target quaternions.
    Args:
        predictions (torch.Tensor): predicted quaternions of shape (batch_size, 4)
        targets (torch.Tensor): target quaternions of shape (batch_size, 4)
    Returns:
        quaternion_loss (torch.Tensor): quaternion loss value
    """
    # Normalize the predictions and targets
    predictions_norm = predictions / torch.linalg.norm(predictions, dim=1, keepdim=True)
    targets_norm = targets / torch.linalg.norm(targets, dim=1, keepdim=True)
    
    # Compute the dot product between the normalized predictions and targets
    dot_product = torch.sum(predictions_norm * targets_norm, dim=1)
    
    # Compute the angle between the predictions and targets
    angle = 2 * torch.acos(torch.clamp(dot_product, min=-1, max=1))
    
    # Compute the quaternion loss
    quaternion_loss = torch.mean(angle)
    
    return quaternion_loss
//  [python]
def calc_str_loss(pred, true, str_mask, A=20.0, gamma=0.95):
    '''
    Calculate structural loss described in the AlphaFold2 patent
    Input:
        - pred: predicted coordinates (I, B, L, n_atom, 3)
        - true: true coordinates (B, L, n_atom, 3)
        - mask: mask for calculating loss (B, L, n_atom, 3)
    Output: str loss
    '''
    I = pred.shape[0]
    true = true.unsqueeze(0)
    t_tilde_ij = get_t(true[:,:,:,0], true[:,:,:,1], true[:,:,:,2])
    t_ij = get_t(pred[:,:,:,0], pred[:,:,:,1], pred[:,:,:,2])
    
    difference = torch.norm(t_tilde_ij-t_ij,dim=-1) # (I, B, L, L)
    temp_mask = torch.full_like(difference, 0, device=pred.device)
    temp_mask[:,:,:,str_mask[0,:]] = 1
    temp_mask[:,:,str_mask[0,:],:] = 1
    loss = -(torch.nn.functional.relu(1.0-difference/A)) * temp_mask ### check why I is 10
    loss = loss.sum(dim=-1).sum(dim=-1).sum(dim=-1)
    mask_sum = temp_mask.sum(dim=-1).sum(dim=-1).sum(dim=-1)
    loss = loss / mask_sum
    # weighting loss
    w_loss = torch.pow(torch.full((I,), gamma, device=pred.device), torch.arange(I, device=pred.device))
    w_loss = torch.flip(w_loss, (0,))
    w_loss = w_loss / w_loss.sum()

    tot_loss = (w_loss * loss).sum()
    return tot_loss, loss.detach()
//  [python]
def calc_str_loss(pred, true, str_mask, A=20.0, gamma=0.95):
    '''
    Calculate structural loss described in the AlphaFold2 patent
    Input:
        - pred: predicted coordinates (I, B, L, n_atom, 3)
        - true: true coordinates (B, L, n_atom, 3)
        - mask: mask for calculating loss (B, L, n_atom, 3)
    Output: str loss
    '''
    I = pred.shape[0]
    true = true.unsqueeze(0)
    t_tilde_ij = get_t(true[:,:,:,0], true[:,:,:,1], true[:,:,:,2])
    t_ij = get_t(pred[:,:,:,0], pred[:,:,:,1], pred[:,:,:,2])
    
    difference = torch.linalg.norm(t_tilde_ij-t_ij,dim=-1) # (I, B, L, L)
    temp_mask = torch.full_like(difference, 0, device=pred.device)
    temp_mask[:,:,:,str_mask[0,:]] = 1
    temp_mask[:,:,str_mask[0,:],:] = 1
    loss = -(torch.nn.functional.relu(1.0-difference/A)) * temp_mask ### check why I is 10
    loss = loss.sum(dim=-1).sum(dim=-1).sum(dim=-1)
    mask_sum = temp_mask.sum(dim=-1).sum(dim=-1).sum(dim=-1)
    loss = loss / mask_sum
    # weighting loss
    w_loss = torch.pow(torch.full((I,), gamma, device=pred.device), torch.arange(I, device=pred.device))
    w_loss = torch.flip(w_loss, (0,))
    w_loss = w_loss / w_loss.sum()

    tot_loss = (w_loss * loss).sum()
    return tot_loss, loss.detach()
//  [python]
def torsion(a,b,c,d, eps=1e-6):
    #A function that takes in 4 atom coordinates:
    # a - [B,L,3]
    # b - [B,L,3]
    # c - [B,L,3]
    # d - [B,L,3]
    # and returns cos and sin of the dihedral angle between those 4 points in order a, b, c, d
    # output - [B,L,2]
    u1 = b-a
    u1 = u1 / (torch.norm(u1, dim=-1, keepdim=True) + eps)
    u2 = c-b
    u2 = u2 / (torch.norm(u2, dim=-1, keepdim=True) + eps)
    u3 = d-c
    u3 = u3 / (torch.norm(u3, dim=-1, keepdim=True) + eps)
    #
    t1 = torch.cross(u1, u2, dim=-1) #[B, L, 3]
    t2 = torch.cross(u2, u3, dim=-1)
    t1_norm = torch.norm(t1, dim=-1, keepdim=True)
    t2_norm = torch.norm(t2, dim=-1, keepdim=True)
    
    cos_angle = torch.matmul(t1[:,:,None,:], t2[:,:,:,None])[:,:,0]
    sin_angle = torch.norm(u2, dim=-1,keepdim=True)*(torch.matmul(u1[:,:,None,:], t2[:,:,:,None])[:,:,0])
    
    cos_sin = torch.cat([cos_angle, sin_angle], axis=-1)/(t1_norm*t2_norm+eps) #[B,L,2]
    return cos_sin
//  [python]
def torsion(a,b,c,d, eps=1e-6):
    #A function that takes in 4 atom coordinates:
    # a - [B,L,3]
    # b - [B,L,3]
    # c - [B,L,3]
    # d - [B,L,3]
    # and returns cos and sin of the dihedral angle between those 4 points in order a, b, c, d
    # output - [B,L,2]
    u1 = b-a
    u1 = u1 / (torch.linalg.norm(u1, dim=-1, keepdim=True) + eps)
    u2 = c-b
    u2 = u2 / (torch.linalg.norm(u2, dim=-1, keepdim=True) + eps)
    u3 = d-c
    u3 = u3 / (torch.linalg.norm(u3, dim=-1, keepdim=True) + eps)
    #
    t1 = torch.cross(u1, u2, dim=-1) #[B, L, 3]
    t2 = torch.cross(u2, u3, dim=-1)
    t1_norm = torch.linalg.norm(t1, dim=-1, keepdim=True)
    t2_norm = torch.linalg.norm(t2, dim=-1, keepdim=True)
    
    cos_angle = torch.matmul(t1[:,:,None,:], t2[:,:,:,None])[:,:,0]
    sin_angle = torch.linalg.norm(u2, dim=-1,keepdim=True)*(torch.matmul(u1[:,:,None,:], t2[:,:,:,None])[:,:,0])
    
    cos_sin = torch.cat([cos_angle, sin_angle], axis=-1)/(t1_norm*t2_norm+eps) #[B,L,2]
    return cos_sin
//  [python]
def calc_temporal_shape_consistency_loss(pred_betas, sequence_inds, weights=None):
    temp_shape_consist_loss = 0

    seq_losses = []
    for seq_inds in sequence_inds:
        seq_pred_betas = pred_betas[seq_inds]
        average_shape = seq_pred_betas.mean(0).unsqueeze(0).detach()
        diff = seq_pred_betas - average_shape
        if weights is not None:
            diff = diff * weights.unsqueeze(0).to(seq_pred_betas.device)
        #(seq_pred_betas.unsqueeze(1) - seq_pred_betas.unsqueeze(0)) * weights.unsqueeze(0).unsqueeze(0).to(seq_pred_betas.device)
        seq_losses.append(torch.norm(diff, p=2, dim=1))
    if len(seq_losses)>0:
        temp_shape_consist_loss = torch.cat(seq_losses, 0)

    return temp_shape_consist_loss
//  [python]
def calc_temporal_shape_consistency_loss(pred_betas, sequence_inds, weights=None):
    temp_shape_consist_loss = 0

    seq_losses = []
    for seq_inds in sequence_inds:
        seq_pred_betas = pred_betas[seq_inds]
        average_shape = seq_pred_betas.mean(0).unsqueeze(0).detach()
        diff = seq_pred_betas - average_shape
        if weights is not None:
            diff = diff * weights.unsqueeze(0).to(seq_pred_betas.device)
        #(seq_pred_betas.unsqueeze(1) - seq_pred_betas.unsqueeze(0)) * weights.unsqueeze(0).unsqueeze(0).to(seq_pred_betas.device)
        seq_losses.append(torch.linalg.norm(diff, ord=2, dim=1))
    if len(seq_losses)>0:
        temp_shape_consist_loss = torch.cat(seq_losses, 0)

    return temp_shape_consist_loss
//  [python]
def _calc_world_trans_loss_(preds, gts, vmasks, sequence_inds):
    # the world trans is supposed to be in shape (B, N, 3), B - batch size, N - temporal length (frame number in a video clip)
    loss = []
    device = preds.device
    gts, vmasks = gts.to(device), vmasks.to(device)
    for seq_inds in sequence_inds:
        pred, gt, vmask = preds[seq_inds], gts[seq_inds], vmasks[seq_inds]
        if vmask.sum()==0:
            continue

        N = pred.shape[0]
        first_index, second_index = clip_frame_pairs_indes(N, device)
        
        scale2align = kabsch_umeyama(gt, pred).detach().clamp(max=10.0)
        #print('scale2align:',scale2align)
        pred_aligned = pred * scale2align

        delta_pred_aligned = pred_aligned[first_index] - pred_aligned[second_index]
        delta_gt = gt[first_index] - gt[second_index]
        error = torch.norm(delta_pred_aligned - delta_gt, dim=-1).mean()

        if torch.isnan(error):
            #print('_calc_world_trans_loss_ is nan', delta_pred_aligned)
            continue
        loss.append(error)
    loss = torch.stack(loss) if len(loss)>0 else torch.zeros(1,device=device)
    return loss
//  [python]
def _calc_world_trans_loss_(preds, gts, vmasks, sequence_inds):
    # the world trans is supposed to be in shape (B, N, 3), B - batch size, N - temporal length (frame number in a video clip)
    loss = []
    device = preds.device
    gts, vmasks = gts.to(device), vmasks.to(device)
    for seq_inds in sequence_inds:
        pred, gt, vmask = preds[seq_inds], gts[seq_inds], vmasks[seq_inds]
        if vmask.sum()==0:
            continue

        N = pred.shape[0]
        first_index, second_index = clip_frame_pairs_indes(N, device)
        
        scale2align = kabsch_umeyama(gt, pred).detach().clamp(max=10.0)
        #print('scale2align:',scale2align)
        pred_aligned = pred * scale2align

        delta_pred_aligned = pred_aligned[first_index] - pred_aligned[second_index]
        delta_gt = gt[first_index] - gt[second_index]
        error = torch.linalg.norm(delta_pred_aligned - delta_gt, dim=-1).mean()

        if torch.isnan(error):
            #print('_calc_world_trans_loss_ is nan', delta_pred_aligned)
            continue
        loss.append(error)
    loss = torch.stack(loss) if len(loss)>0 else torch.zeros(1,device=device)
    return loss
//  [python]
def u(self):
        u_unnormed = self.u_unnormed() if callable(self.u_unnormed) else self.u_unnormed
        norm = torch.norm(u_unnormed, p=2, dim=-1, keepdim=True)
        return torch.div(u_unnormed, norm)
//  [python]
def u(self):
        u_unnormed = self.u_unnormed() if callable(self.u_unnormed) else self.u_unnormed
        norm = torch.linalg.norm(u_unnormed, ord=2, dim=-1, keepdim=True)
        return torch.div(u_unnormed, norm)
//  [python]
def forward(self, verts_pred, verts_gt, smooth=False):
        verts_pred = torch.matmul(self.ls, verts_pred)
        verts_gt = torch.matmul(self.ls, verts_gt)
        loss = torch.norm(verts_pred - verts_gt, dim=-1).mean()
        if smooth:
            loss_smooth = torch.norm(torch.matmul(self.ls, verts_pred), dim=-1).mean()
            return loss, loss_smooth
        return loss, None
//  [python]
def forward(self, verts_pred, verts_gt, smooth=False):
        verts_pred = torch.matmul(self.ls, verts_pred)
        verts_gt = torch.matmul(self.ls, verts_gt)
        loss = torch.linalg.norm(verts_pred - verts_gt, dim=-1).mean()
        if smooth:
            loss_smooth = torch.linalg.norm(torch.matmul(self.ls, verts_pred), dim=-1).mean()
            return loss, loss_smooth
        return loss, None
//  [python]
def calc_vel_acc_error(seq_gts, seq_preds, seq_mask):
    velocity_mask = (seq_mask[1:]*seq_mask[:-1]).float()
    vel_gts = seq_gts[1:] - seq_gts[:-1]
    vel_preds = seq_preds[1:] - seq_preds[:-1]
    velocity_loss = (vel_gts - vel_preds) * velocity_mask[:,None]
    velocity_loss = torch.norm(velocity_loss, p=2, dim=1)

    acc_mask = (velocity_mask[1:]*velocity_mask[:-1]).float()
    acc_gts = vel_gts[1:] - vel_gts[:-1]
    acc_preds = vel_preds[1:] - vel_preds[:-1]
    acc_loss = (acc_gts - acc_preds) * acc_mask[:,None]
    acc_loss = torch.norm(acc_loss, p=2, dim=1)
    # to be equal length with the vel loss
    acc_loss = torch.cat([torch.zeros_like(acc_loss)[[0]], acc_loss],0)

    return velocity_loss, acc_loss
//  [python]
def calc_vel_acc_error(seq_gts, seq_preds, seq_mask):
    velocity_mask = (seq_mask[1:]*seq_mask[:-1]).float()
    vel_gts = seq_gts[1:] - seq_gts[:-1]
    vel_preds = seq_preds[1:] - seq_preds[:-1]
    velocity_loss = (vel_gts - vel_preds) * velocity_mask[:,None]
    velocity_loss = torch.linalg.norm(velocity_loss, p=2, dim=1)

    acc_mask = (velocity_mask[1:]*velocity_mask[:-1]).float()
    acc_gts = vel_gts[1:] - vel_gts[:-1]
    acc_preds = vel_preds[1:] - vel_preds[:-1]
    acc_loss = (acc_gts - acc_preds) * acc_mask[:,None]
    acc_loss = torch.linalg.norm(acc_loss, p=2, dim=1)
    # to be equal length with the vel loss
    acc_loss = torch.cat([torch.zeros_like(acc_loss)[[0]], acc_loss],0)

    return velocity_loss, acc_loss
//  [python]
def calc_motion_offsets3D_loss(motion_offsets, clip_frame_ids, traj3D_gts):
    if len(traj3D_gts.shape) == 3:
        valid_3Doffset_masks, offset3D_gts = get_valid_offset_mask(traj3D_gts, clip_frame_ids)
    elif len(traj3D_gts.shape) == 2:
        valid_3Doffset_masks, offset3D_gts = get_valid_offset_maskV2(traj3D_gts, clip_frame_ids)
        
    if valid_3Doffset_masks.sum()>0:
        valid_3Doffset_masks = valid_3Doffset_masks.detach()
        offset3D_loss = torch.norm(motion_offsets[valid_3Doffset_masks] - offset3D_gts.to(motion_offsets.device), p=2, dim=-1)
    else:
        offset3D_loss = 0
    return offset3D_loss
//  [python]
def calc_motion_offsets3D_loss(motion_offsets, clip_frame_ids, traj3D_gts):
    if len(traj3D_gts.shape) == 3:
        valid_3Doffset_masks, offset3D_gts = get_valid_offset_mask(traj3D_gts, clip_frame_ids)
    elif len(traj3D_gts.shape) == 2:
        valid_3Doffset_masks, offset3D_gts = get_valid_offset_maskV2(traj3D_gts, clip_frame_ids)
        
    if valid_3Doffset_masks.sum()>0:
        valid_3Doffset_masks = valid_3Doffset_masks.detach()
        offset3D_loss = torch.linalg.norm(motion_offsets[valid_3Doffset_masks] - offset3D_gts.to(motion_offsets.device), ord=2, dim=-1)
    else:
        offset3D_loss = 0
    return offset3D_loss
//  [python]
def calc_cam_loss_with_full_body_bboxes(all_pred_cams, all_pred_verts, full_body_bboxes, focal_length):
    valid_mask = (full_body_bboxes!=-2).sum(-1) == 0
    if valid_mask.sum() == 0:
        return 0
    pred_verts = all_pred_verts[valid_mask]
    fb_bboxes = full_body_bboxes[valid_mask]
    
    bbox3D = get_3Dbbox(pred_verts)

    if args().estimate_camera:
        pred_trans = all_pred_cams[valid_mask]
        pred_fovs = focal_length[valid_mask]
        bbox_2Dprojection = perspective_projection_withfovs(
                        bbox3D, translation=pred_trans, fovs=pred_fovs)
    else:
        pred_cams = all_pred_cams[valid_mask]
        trans_preds = denormalize_cam_params_to_trans(pred_cams, positive_constrain=False).to(bbox3D.device)
        bbox_2Dprojection = perspective_projection(bbox3D, translation=trans_preds, focal_length=focal_length, normalize=True)[:,:,:2].float()
    
    bbox3D_2Dcenter_projection_preds = (bbox_2Dprojection.min(1).values + bbox_2Dprojection.max(1).values) / 2
    bbox2Dcenter_gts = (fb_bboxes[:, :2] + fb_bboxes[:, 2:]) / 2
    
    fbbox_cam_loss = torch.norm(bbox3D_2Dcenter_projection_preds - bbox2Dcenter_gts, p=2, dim=1)
    return fbbox_cam_loss
//  [python]
def calc_cam_loss_with_full_body_bboxes(all_pred_cams, all_pred_verts, full_body_bboxes, focal_length):
    valid_mask = (full_body_bboxes!=-2).sum(-1) == 0
    if valid_mask.sum() == 0:
        return 0
    pred_verts = all_pred_verts[valid_mask]
    fb_bboxes = full_body_bboxes[valid_mask]
    
    bbox3D = get_3Dbbox(pred_verts)

    if args().estimate_camera:
        pred_trans = all_pred_cams[valid_mask]
        pred_fovs = focal_length[valid_mask]
        bbox_2Dprojection = perspective_projection_withfovs(
                        bbox3D, translation=pred_trans, fovs=pred_fovs)
    else:
        pred_cams = all_pred_cams[valid_mask]
        trans_preds = denormalize_cam_params_to_trans(pred_cams, positive_constrain=False).to(bbox3D.device)
        bbox_2Dprojection = perspective_projection(bbox3D, translation=trans_preds, focal_length=focal_length, normalize=True)[:,:,:2].float()
    
    bbox3D_2Dcenter_projection_preds = (bbox_2Dprojection.min(1).values + bbox_2Dprojection.max(1).values) / 2
    bbox2Dcenter_gts = (fb_bboxes[:, :2] + fb_bboxes[:, 2:]) / 2
    
    fbbox_cam_loss = torch.linalg.norm(bbox3D_2Dcenter_projection_preds - bbox2Dcenter_gts, ord=2, dim=1)
    return fbbox_cam_loss
//  [python]
def select_experts(gate_logits: torch.Tensor, top_k: int, moe_normalize_expert_weights: int):
    # all_probs: (sequence_length, n_experts) and upcast for softmax
    all_probs = torch.nn.functional.softmax(gate_logits, dim=1, dtype=torch.float)
    # weights, selected_experts: (sequence_length, top-k)
    weights, selected_experts = torch.topk(all_probs, top_k, dim=-1)
    if moe_normalize_expert_weights:
        weights = weights / torch.norm(weights, p=moe_normalize_expert_weights, dim=-1, keepdim=True)
    weights = weights.view(-1)
    selected_experts = selected_experts.view(-1)

    return selected_experts, weights
//  [python]
def select_experts(gate_logits: torch.Tensor, top_k: int, moe_normalize_expert_weights: int):
    # all_probs: (sequence_length, n_experts) and upcast for softmax
    all_probs = torch.nn.functional.softmax(gate_logits, dim=1, dtype=torch.float)
    # weights, selected_experts: (sequence_length, top-k)
    weights, selected_experts = torch.topk(all_probs, top_k, dim=-1)
    if moe_normalize_expert_weights:
        weights = weights / torch.linalg.norm(weights, ord=moe_normalize_expert_weights, dim=-1, keepdim=True)
    weights = weights.view(-1)
    selected_experts = selected_experts.view(-1)

    return selected_experts, weights
//  [python]
def _unproject_points_batch(
        self,
        focal,
        principal_point,
        radial_params,
        tangential_params,
        thin_prism_params,
        xy: torch.Tensor,
    ) -> torch.Tensor:
        """
        Args:
            focal: (1)
            principal_point: (2)
            radial_params: (num_radial)
            tangential_params: (2)
            thin_prism_params: (4)
            xy: (..., 2)

        Returns:
            point3d_est: (..., 3)
        """
        sh = list(xy.shape[:-1])
        assert xy.shape[-1] == 2, "xy_depth shape incorrect"
        uv_distorted = (xy - principal_point) / focal

        # get xr_yr from uvDistorted
        xr_yr = self._compute_xr_yr_from_uv_distorted(
            tangential_params, thin_prism_params, uv_distorted
        )
        xr_yrNorm = torch.norm(xr_yr, dim=-1)

        # find theta
        theta = self._get_theta_from_norm_xr_yr(radial_params, xr_yrNorm)
        # get the point coordinates:
        point3d_est = theta.new_ones(*sh, 3)
        point3d_est[..., :2] = theta.tan()[..., None] / xr_yrNorm[..., None] * xr_yr
        return point3d_est
//  [python]
def _unproject_points_batch(
        self,
        focal,
        principal_point,
        radial_params,
        tangential_params,
        thin_prism_params,
        xy: torch.Tensor,
    ) -> torch.Tensor:
        """
        Args:
            focal: (1)
            principal_point: (2)
            radial_params: (num_radial)
            tangential_params: (2)
            thin_prism_params: (4)
            xy: (..., 2)

        Returns:
            point3d_est: (..., 3)
        """
        sh = list(xy.shape[:-1])
        assert xy.shape[-1] == 2, "xy_depth shape incorrect"
        uv_distorted = (xy - principal_point) / focal

        # get xr_yr from uvDistorted
        xr_yr = self._compute_xr_yr_from_uv_distorted(
            tangential_params, thin_prism_params, uv_distorted
        )
        xr_yrNorm = torch.linalg.norm(xr_yr, dim=-1)

        # find theta
        theta = self._get_theta_from_norm_xr_yr(radial_params, xr_yrNorm)
        # get the point coordinates:
        point3d_est = theta.new_ones(*sh, 3)
        point3d_est[..., :2] = theta.tan()[..., None] / xr_yrNorm[..., None] * xr_yr
        return point3d_est
//  [python]
def mpjpe_byjoint(predicted, target):
    """
    Mean per-joint position error (i.e. mean Euclidean distance),
    often referred to as "Protocol #1" in many papers.
    """
    assert predicted.shape == target.shape
    return torch.mean(torch.norm(predicted - target, dim=len(target.shape) - 1), dim=0)
//  [python]
def mpjpe_byjoint(predicted, target):
    """
    Mean per-joint position error (i.e. mean Euclidean distance),
    often referred to as "Protocol #1" in many papers.
    """
    assert predicted.shape == target.shape
    return torch.mean(torch.linalg.norm(predicted - target, dim=len(target.shape) - 1), dim=0)
//  [python]
def compute_diagonal_cov_reg(feats) -> Tensor:
    """
    Computes diagonal covariance regularization
    """
    assert feats.ndim == 2, f"The function accepts only 2d feature vectors: {feats.shape}"
    cov = compute_covariance(feats)
    eye = torch.eye(feats.size(1), device=cov.device)

    return torch.norm(cov - eye)
//  [python]
def compute_diagonal_cov_reg(feats) -> Tensor:
    """
    Computes diagonal covariance regularization
    """
    assert feats.ndim == 2, f"The function accepts only 2d feature vectors: {feats.shape}"
    cov = compute_covariance(feats)
    eye = torch.eye(feats.size(1), device=cov.device)

    return torch.linalg.norm(cov - eye)
//  [python]
def extract_bbox_clusters(
    bbox_features_file: str,
    output_file: str,
    num_clusters: int = 20, 
    seed: int = 0, 
    pca_dim: Optional[int] = 0,
):
    """
    Example:
        python extract.py extract_bbox_clusters \
            --bbox_features_file "./data/VOC2012/multi_region_bboxes/fixed/bbox_features_e2_d5.pth" \
            --pca_dim 32 --num_clusters 21 --seed 0 \
            --output_file "./data/VOC2012/multi_region_bboxes/fixed/bbox_clusters_e2_d5_pca_32.pth" \
    """

    # Load bounding boxes
    bbox_list = torch.load(bbox_features_file)
    total_num_boxes = sum(len(d['bboxes']) for d in bbox_list)
    print(f'Loaded bounding box list. There are {total_num_boxes} total bounding boxes with features.')

    # Loop over boxes and stack features with PyTorch, because Numpy is too slow
    print(f'Stacking and normalizing features')
    all_features = torch.cat([bbox_dict['features'] for bbox_dict in bbox_list], dim=0)  # (numBbox, D)
    all_features = all_features / torch.norm(all_features, dim=-1, keepdim=True)  # (numBbox, D)f
    all_features = all_features.numpy()

    # Cluster: PCA
    if pca_dim:
        pca = PCA(pca_dim)
        print(f'Computing PCA with dimension {pca_dim}')
        all_features = pca.fit_transform(all_features)

    # Cluster: K-Means
    print(f'Computing K-Means clustering with {num_clusters} clusters')
    kmeans = MiniBatchKMeans(n_clusters=num_clusters, batch_size=4096, max_iter=5000, random_state=seed)
    clusters = kmeans.fit_predict(all_features)
    
    # Print 
    _indices, _counts = np.unique(clusters, return_counts=True)
    print(f'Cluster indices: {_indices.tolist()}')
    print(f'Cluster counts: {_counts.tolist()}')

    # Loop over boxes and add clusters
    idx = 0
    for bbox_dict in bbox_list:
        num_bboxes = len(bbox_dict['bboxes'])
        del bbox_dict['features']  # bbox_dict['features'] = bbox_dict['features'].squeeze()
        bbox_dict['clusters'] = clusters[idx: idx + num_bboxes]
        idx = idx + num_bboxes
    
    # Save
    torch.save(bbox_list, output_file)
    print(f'Saved features to {output_file}')
//  [python]
def extract_bbox_clusters(
    bbox_features_file: str,
    output_file: str,
    num_clusters: int = 20, 
    seed: int = 0, 
    pca_dim: Optional[int] = 0,
):
    """
    Example:
        python extract.py extract_bbox_clusters \
            --bbox_features_file "./data/VOC2012/multi_region_bboxes/fixed/bbox_features_e2_d5.pth" \
            --pca_dim 32 --num_clusters 21 --seed 0 \
            --output_file "./data/VOC2012/multi_region_bboxes/fixed/bbox_clusters_e2_d5_pca_32.pth" \
    """

    # Load bounding boxes
    bbox_list = torch.load(bbox_features_file)
    total_num_boxes = sum(len(d['bboxes']) for d in bbox_list)
    print(f'Loaded bounding box list. There are {total_num_boxes} total bounding boxes with features.')

    # Loop over boxes and stack features with PyTorch, because Numpy is too slow
    print(f'Stacking and normalizing features')
    all_features = torch.cat([bbox_dict['features'] for bbox_dict in bbox_list], dim=0)  # (numBbox, D)
    all_features = all_features / torch.linalg.norm(all_features, dim=-1, keepdim=True)  # (numBbox, D)f
    all_features = all_features.numpy()

    # Cluster: PCA
    if pca_dim:
        pca = PCA(pca_dim)
        print(f'Computing PCA with dimension {pca_dim}')
        all_features = pca.fit_transform(all_features)

    # Cluster: K-Means
    print(f'Computing K-Means clustering with {num_clusters} clusters')
    kmeans = MiniBatchKMeans(n_clusters=num_clusters, batch_size=4096, max_iter=5000, random_state=seed)
    clusters = kmeans.fit_predict(all_features)
    
    # Print 
    _indices, _counts = np.unique(clusters, return_counts=True)
    print(f'Cluster indices: {_indices.tolist()}')
    print(f'Cluster counts: {_counts.tolist()}')

    # Loop over boxes and add clusters
    idx = 0
    for bbox_dict in bbox_list:
        num_bboxes = len(bbox_dict['bboxes'])
        del bbox_dict['features']  # bbox_dict['features'] = bbox_dict['features'].squeeze()
        bbox_dict['clusters'] = clusters[idx: idx + num_bboxes]
        idx = idx + num_bboxes
    
    # Save
    torch.save(bbox_list, output_file)
    print(f'Saved features to {output_file}')
//  [python]
def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        x: (sequence_length, model_dim)
        gate_logits: (sequence_length, n_experts)
        """
        # optional reshape
        input_shape = x.shape
        x = x.view(-1, input_shape[-1])

        # gate_logits: (sequence_length, n_experts)
        gate_logits = self.gate(x)
        # all_probs: (sequence_length, n_experts) and upcast for softmax
        weights = torch.nn.functional.softmax(gate_logits, dim=1, dtype=torch.float)

        if self.top_k < self.num_experts:
            _, not_selected_experts = torch.topk(
                weights,
                self.num_experts - self.top_k,
                largest=False,
                sorted=False,
                dim=1,
            )
            # Mask not selected experts
            weights.scatter_(1, not_selected_experts, 0)

        # Re-normalize
        if self.moe_normalize_expert_weights:
            weights = weights / torch.norm(weights, p=self.moe_normalize_expert_weights, dim=-1, keepdim=True)
        weights = weights.to(x.dtype)

        # Final output tensor
        out = x.new_zeros(x.shape[0], self.hidden_dim)
        for i in range(self.num_experts):
            h = self.act(self.w1[i](x)) * self.v1[i](x)
            h = self.w2[i](h, reduce=False)
            # Add expert output to out with masking
            out += h * weights[:, i].view(-1, 1)

        # Reduce sum
        if self.process_group.size() > 1:
            torch.distributed.all_reduce(out, group=self.process_group)

        return out
//  [python]
def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        x: (sequence_length, model_dim)
        gate_logits: (sequence_length, n_experts)
        """
        # optional reshape
        input_shape = x.shape
        x = x.view(-1, input_shape[-1])

        # gate_logits: (sequence_length, n_experts)
        gate_logits = self.gate(x)
        # all_probs: (sequence_length, n_experts) and upcast for softmax
        weights = torch.nn.functional.softmax(gate_logits, dim=1, dtype=torch.float)

        if self.top_k < self.num_experts:
            _, not_selected_experts = torch.topk(
                weights,
                self.num_experts - self.top_k,
                largest=False,
                sorted=False,
                dim=1,
            )
            # Mask not selected experts
            weights.scatter_(1, not_selected_experts, 0)

        # Re-normalize
        if self.moe_normalize_expert_weights:
            weights = weights / torch.linalg.norm(weights, ord=self.moe_normalize_expert_weights, dim=-1, keepdim=True)
        weights = weights.to(x.dtype)

        # Final output tensor
        out = x.new_zeros(x.shape[0], self.hidden_dim)
        for i in range(self.num_experts):
            h = self.act(self.w1[i](x)) * self.v1[i](x)
            h = self.w2[i](h, reduce=False)
            # Add expert output to out with masking
            out += h * weights[:, i].view(-1, 1)

        # Reduce sum
        if self.process_group.size() > 1:
            torch.distributed.all_reduce(out, group=self.process_group)

        return out
//  [python]
def dense_forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        x: (sequence_length, model_dim)
        gate_logits: (sequence_length, n_experts)
        """
        # optional reshape
        input_shape = x.shape
        x = x.view(-1, input_shape[-1])

        # gate_logits: (sequence_length, n_experts)
        gate_logits = self.gate(x)
        # all_probs: (sequence_length, n_experts) and upcast for softmax
        weights = torch.nn.functional.softmax(gate_logits, dim=1, dtype=torch.float)

        if self.top_k < self.num_experts:
            _, not_selected_experts = torch.topk(
                weights,
                self.num_experts - self.top_k,
                largest=False,
                sorted=False,
                dim=1,
            )
            # Mask not selected experts
            weights.scatter_(1, not_selected_experts, 0)

        # Re-normalize
        if self.moe_normalize_expert_weights:
            weights = weights / torch.norm(weights, p=self.moe_normalize_expert_weights, dim=-1, keepdim=True)
        weights = weights.to(x.dtype)

        # Expand to [num_experts, sequence_length, model_dim]
        x = x.view(1, -1, input_shape[-1]).expand(self.num_experts, -1, input_shape[-1])

        # Permute to [num_experts, model_dim, ffn_dim]
        w1 = self.w1.view(self.num_experts, self.ffn_dim, self.hidden_dim).permute(0, 2, 1)
        v1 = self.v1.view(self.num_experts, self.ffn_dim, self.hidden_dim).permute(0, 2, 1)

        inter = self.act(torch.bmm(x, w1)) * torch.bmm(x, v1)

        out = torch.bmm(inter, self.w2.view(self.num_experts, self.ffn_dim, self.hidden_dim))
        # Mask not selected experts
        out *= weights.t().view(self.num_experts, -1, 1)

        # Sum experts
        out = out.sum(0)

        # Reduce sum
        if self.process_group.size() > 1:
            torch.distributed.all_reduce(out, group=self.process_group)

        return out
//  [python]
def dense_forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        x: (sequence_length, model_dim)
        gate_logits: (sequence_length, n_experts)
        """
        # optional reshape
        input_shape = x.shape
        x = x.view(-1, input_shape[-1])

        # gate_logits: (sequence_length, n_experts)
        gate_logits = self.gate(x)
        # all_probs: (sequence_length, n_experts) and upcast for softmax
        weights = torch.nn.functional.softmax(gate_logits, dim=1, dtype=torch.float)

        if self.top_k < self.num_experts:
            _, not_selected_experts = torch.topk(
                weights,
                self.num_experts - self.top_k,
                largest=False,
                sorted=False,
                dim=1,
            )
            # Mask not selected experts
            weights.scatter_(1, not_selected_experts, 0)

        # Re-normalize
        if self.moe_normalize_expert_weights:
            weights = weights / torch.linalg.norm(weights, ord=self.moe_normalize_expert_weights, dim=-1, keepdim=True)
        weights = weights.to(x.dtype)

        # Expand to [num_experts, sequence_length, model_dim]
        x = x.view(1, -1, input_shape[-1]).expand(self.num_experts, -1, input_shape[-1])

        # Permute to [num_experts, model_dim, ffn_dim]
        w1 = self.w1.view(self.num_experts, self.ffn_dim, self.hidden_dim).permute(0, 2, 1)
        v1 = self.v1.view(self.num_experts, self.ffn_dim, self.hidden_dim).permute(0, 2, 1)

        inter = self.act(torch.bmm(x, w1)) * torch.bmm(x, v1)

        out = torch.bmm(inter, self.w2.view(self.num_experts, self.ffn_dim, self.hidden_dim))
        # Mask not selected experts
        out *= weights.t().view(self.num_experts, -1, 1)

        # Sum experts
        out = out.sum(0)

        # Reduce sum
        if self.process_group.size() > 1:
            torch.distributed.all_reduce(out, group=self.process_group)

        return out
//  [python]
def get_weighted_loss(
        self,
        losses,
        shared_parameters,
        **kwargs,
    ):
        """

        Parameters
        ----------
        losses :
        shared_parameters : shared parameters
        kwargs :

        Returns
        -------

        """

        extra_outputs = dict()
        if self.step == 0:
            self._init_optim_problem()

        if (self.step % self.update_weights_every) == 0:
            self.step += 1

            grads = {}
            for i, loss in enumerate(losses):
                g = list(
                    torch.autograd.grad(
                        loss,
                        shared_parameters,
                        retain_graph=True,
                    )
                )
                grad = torch.cat([torch.flatten(grad) for grad in g])
                grads[i] = grad

            G = torch.stack(tuple(v for v in grads.values()))
            GTG = torch.mm(G, G.t())

            self.normalization_factor = (
                torch.norm(GTG).detach().cpu().numpy().reshape((1,))
            )
            GTG = GTG / self.normalization_factor.item()
            alpha = self.solve_optimization(GTG.cpu().detach().numpy())
            alpha = torch.from_numpy(alpha)

        else:
            self.step += 1
            alpha = self.prvs_alpha

        weighted_loss = sum([losses[i] * alpha[i] for i in range(len(alpha))])
        extra_outputs["weights"] = alpha
        return weighted_loss, extra_outputs
//  [python]
def get_weighted_loss(
        self,
        losses,
        shared_parameters,
        **kwargs,
    ):
        """

        Parameters
        ----------
        losses :
        shared_parameters : shared parameters
        kwargs :

        Returns
        -------

        """

        extra_outputs = dict()
        if self.step == 0:
            self._init_optim_problem()

        if (self.step % self.update_weights_every) == 0:
            self.step += 1

            grads = {}
            for i, loss in enumerate(losses):
                g = list(
                    torch.autograd.grad(
                        loss,
                        shared_parameters,
                        retain_graph=True,
                    )
                )
                grad = torch.cat([torch.flatten(grad) for grad in g])
                grads[i] = grad

            G = torch.stack(tuple(v for v in grads.values()))
            GTG = torch.mm(G, G.t())

            self.normalization_factor = (
                torch.linalg.norm(GTG).detach().cpu().numpy().reshape((1,))
            )
            GTG = GTG / self.normalization_factor.item()
            alpha = self.solve_optimization(GTG.cpu().detach().numpy())
            alpha = torch.from_numpy(alpha)

        else:
            self.step += 1
            alpha = self.prvs_alpha

        weighted_loss = sum([losses[i] * alpha[i] for i in range(len(alpha))])
        extra_outputs["weights"] = alpha
        return weighted_loss, extra_outputs
//  [python]
def _project_conflicting(self, grads: List[Tuple[torch.Tensor]]):
        pc_grad = copy.deepcopy(grads)
        for g_i in pc_grad:
            random.shuffle(grads)
            for g_j in grads:
                g_i_g_j = sum(
                    [
                        torch.dot(torch.flatten(grad_i), torch.flatten(grad_j))
                        for grad_i, grad_j in zip(g_i, g_j)
                    ]
                )
                if g_i_g_j < 0:
                    g_j_norm_square = (
                        torch.norm(torch.cat([torch.flatten(g) for g in g_j])) ** 2
                    )
                    for grad_i, grad_j in zip(g_i, g_j):
                        grad_i -= g_i_g_j * grad_j / g_j_norm_square

        merged_grad = [sum(g) for g in zip(*pc_grad)]
        if self.reduction == "mean":
            merged_grad = [g / self.n_tasks for g in merged_grad]

        return merged_grad
//  [python]
def _project_conflicting(self, grads: List[Tuple[torch.Tensor]]):
        pc_grad = copy.deepcopy(grads)
        for g_i in pc_grad:
            random.shuffle(grads)
            for g_j in grads:
                g_i_g_j = sum(
                    [
                        torch.dot(torch.flatten(grad_i), torch.flatten(grad_j))
                        for grad_i, grad_j in zip(g_i, g_j)
                    ]
                )
                if g_i_g_j < 0:
                    g_j_norm_square = (
                        torch.linalg.norm(torch.cat([torch.flatten(g) for g in g_j])) ** 2
                    )
                    for grad_i, grad_j in zip(g_i, g_j):
                        grad_i -= g_i_g_j * grad_j / g_j_norm_square

        merged_grad = [sum(g) for g in zip(*pc_grad)]
        if self.reduction == "mean":
            merged_grad = [g / self.n_tasks for g in merged_grad]

        return merged_grad