jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def _compute_norm(t, n, dim):
    r"""Compute the L_n-norm of a tensor along all dimensions except for the specified dimension.

    The L_n-norm will be computed across all entries in tensor `t` along all dimension
    except for the one identified by dim.
    Example: if `t` is of shape, say, 3x2x4 and dim=2 (the last dim),
    then norm will have Size [4], and each entry will represent the
    `L_n`-norm computed using the 3x2=6 entries for each of the 4 channels.

    Args:
        t (torch.Tensor): tensor representing the parameter to prune
        n (int, float, inf, -inf, 'fro', 'nuc'): See documentation of valid
            entries for argument p in torch.norm
        dim (int): dim identifying the channels to prune

    Returns:
        norm (torch.Tensor): L_n norm computed across all dimensions except
            for `dim`. By construction, `norm.shape = t.shape[-1]`.
    """
    # dims = all axes, except for the one identified by `dim`
    dims = list(range(t.dim()))
    # convert negative indexing
    if dim < 0:
        dim = dims[dim]
    dims.remove(dim)

    norm = torch.linalg.norm(t, ord=n, dim=dims)
    return norm
//  [python]
def loss_reduction(sample_loss, reduction):
    if reduction == "mean":
        return torch.mean(sample_loss)
    elif reduction == "max":
        return torch.max(sample_loss)
    elif reduction == "4norm":
        return torch.norm(sample_loss, p=4)
    else:
        raise Exception("unknown reduction method")
//  [python]
def loss_reduction(sample_loss, reduction):
    if reduction == "mean":
        return torch.mean(sample_loss)
    elif reduction == "max":
        return torch.max(sample_loss)
    elif reduction == "4norm":
        return torch.linalg.norm(sample_loss, p=4)
    else:
        raise Exception("unknown reduction method")
//  [python]
def get_rotation_matrix(v, eps=10e-5):
    v = v / (torch.norm(v, dim=-1, keepdim=True) + eps)
    rot = torch.stack((
        torch.stack((v[:, 0], v[:, 1]), dim=-1),
        torch.stack((-v[:, 1], v[:, 0]), dim=-1),
        torch.zeros(v.size(0), 2).type_as(v)
    ), dim=-1)
    return rot
//  [python]
def get_rotation_matrix(v, eps=10e-5):
    v = v / (torch.linalg.norm(v, dim=-1, keepdim=True) + eps)
    rot = torch.stack((
        torch.stack((v[:, 0], v[:, 1]), dim=-1),
        torch.stack((-v[:, 1], v[:, 0]), dim=-1),
        torch.zeros(v.size(0), 2).type_as(v)
    ), dim=-1)
    return rot
//  [python]
def J_normalization(x):
    """Jacobian of the L2 normalization, assuming that we normalize along the
    last dimension."""
    x_normed = torch.nn.functional.normalize(x, dim=-1)
    norm = torch.norm(x, p=2, dim=-1, keepdim=True)

    Id = torch.diag_embed(torch.ones_like(x_normed))
    J = Id - x_normed.unsqueeze(-1) @ x_normed.unsqueeze(-2)
    J = J / norm.unsqueeze(-1)
    return J
//  [python]
def J_normalization(x):
    """Jacobian of the L2 normalization, assuming that we normalize along the
    last dimension."""
    x_normed = torch.nn.functional.normalize(x, dim=-1)
    norm = torch.linalg.norm(x, p=2, dim=-1, keepdim=True)

    Id = torch.diag_embed(torch.ones_like(x_normed))
    J = Id - x_normed.unsqueeze(-1) @ x_normed.unsqueeze(-2)
    J = J / norm.unsqueeze(-1)
    return J
//  [python]
def expm_krylov(A, x, phi=1, krylov_dim=30, inplace=True):

    if krylov_dim > x.size()[0]:
        krylov_dim = x.size()[0]

    Q = th.zeros(x.size()[0], krylov_dim, dtype=x.dtype, device=x.device)
    T = th.zeros(krylov_dim + 1, krylov_dim + 1, dtype=x.dtype, device=x.device)

    #compute the norm of the vector
    norm_x = th.norm(x, p=2)

    # normalize vector
    q = x/norm_x
    Q[:, 0] = q.clone()

    r = band_mv(A, q)

    T[0, 0] = th.dot(q, r)
    r = r - q.mul(T[0, 0])


    T[0, 1] = th.norm(r, p=2) + 1e-10
    T[1, 0] = T[0, 1]

    for k in range(1, krylov_dim):
        b = q
        q = r

        q.div_(T[k -1, k])

        Q[:,k] = q.clone()

        r = band_mv(A, q) - b.mul_(T[k -1, k])

        T[k, k] = th.dot(q, r)

        r =  r - q.mul(T[k, k])

        T[k + 1, k] = th.norm(r, p=2) + 1e-10
        T[k, k + 1] = T[k + 1, k]

    T.mul_(phi)
    exp_mat = expm_eig(T[:-1, :-1])

    if inplace:
        th.mv(Q, exp_mat[:,0], out=x)
        x.mul_(norm_x)
    else:
        return  th.mv(Q, exp_mat[:,0]).mul_(x)
//  [python]
def expm_krylov(A, x, phi=1, krylov_dim=30, inplace=True):

    if krylov_dim > x.size()[0]:
        krylov_dim = x.size()[0]

    Q = th.zeros(x.size()[0], krylov_dim, dtype=x.dtype, device=x.device)
    T = th.zeros(krylov_dim + 1, krylov_dim + 1, dtype=x.dtype, device=x.device)

    #compute the norm of the vector
    norm_x = th.linalg.norm(x, p=2)

    # normalize vector
    q = x/norm_x
    Q[:, 0] = q.clone()

    r = band_mv(A, q)

    T[0, 0] = th.dot(q, r)
    r = r - q.mul(T[0, 0])


    T[0, 1] = th.linalg.norm(r, p=2) + 1e-10
    T[1, 0] = T[0, 1]

    for k in range(1, krylov_dim):
        b = q
        q = r

        q.div_(T[k -1, k])

        Q[:,k] = q.clone()

        r = band_mv(A, q) - b.mul_(T[k -1, k])

        T[k, k] = th.dot(q, r)

        r =  r - q.mul(T[k, k])

        T[k + 1, k] = th.linalg.norm(r, p=2) + 1e-10
        T[k, k + 1] = T[k + 1, k]

    T.mul_(phi)
    exp_mat = expm_eig(T[:-1, :-1])

    if inplace:
        th.mv(Q, exp_mat[:,0], out=x)
        x.mul_(norm_x)
    else:
        return  th.mv(Q, exp_mat[:,0]).mul_(x)
//  [python]
def l_infinity_gradient(x, *, max_tol=0.) -> torch.Tensor:
    """
    Compute the gradient of the infinity-norm |x|∞
    Args:
      max_tol: If |x[i]| is within max_tol to |x|∞, then we consider the
      gradient w.r.t x[i].

    Return:
      gradient: A (num_possible_gradient x x_dim) size tensor. gradient[i] is
      the i'th possible gradient.
    """
    assert (isinstance(x, torch.Tensor))
    x_dim = x.shape[0]
    assert (x.shape == (x_dim, ))
    inf_norm = torch.norm(x, p=float("inf"))
    if inf_norm <= max_tol:
        return torch.cat((torch.eye(
            x_dim, dtype=x.dtype), -torch.eye(x_dim, dtype=x.dtype)),
                         dim=0)
    grad = torch.where(torch.abs(x - inf_norm) <= max_tol, 1, 0) + \
        torch.where(torch.abs(x + inf_norm) <= max_tol, -1, 0)
    nonzero_grad_indices = torch.nonzero(grad).squeeze(1).tolist()
    if len(nonzero_grad_indices) == 1:
        return grad.unsqueeze(0).type(x.dtype)
    all_grad = torch.zeros((len(nonzero_grad_indices), x_dim), dtype=x.dtype)
    for (i, grad_index) in enumerate(nonzero_grad_indices):
        all_grad[i, grad_index] = grad[grad_index]
    return all_grad
//  [python]
def l_infinity_gradient(x, *, max_tol=0.) -> torch.Tensor:
    """
    Compute the gradient of the infinity-norm |x|∞
    Args:
      max_tol: If |x[i]| is within max_tol to |x|∞, then we consider the
      gradient w.r.t x[i].

    Return:
      gradient: A (num_possible_gradient x x_dim) size tensor. gradient[i] is
      the i'th possible gradient.
    """
    assert (isinstance(x, torch.Tensor))
    x_dim = x.shape[0]
    assert (x.shape == (x_dim, ))
    inf_norm = torch.linalg.norm(x, ord=float("inf"))
    if inf_norm <= max_tol:
        return torch.cat((torch.eye(
            x_dim, dtype=x.dtype), -torch.eye(x_dim, dtype=x.dtype)),
                         dim=0)
    grad = torch.where(torch.abs(x - inf_norm) <= max_tol, 1, 0) + \
        torch.where(torch.abs(x + inf_norm) <= max_tol, -1, 0)
    nonzero_grad_indices = torch.nonzero(grad).squeeze(1).tolist()
    if len(nonzero_grad_indices) == 1:
        return grad.unsqueeze(0).type(x.dtype)
    all_grad = torch.zeros((len(nonzero_grad_indices), x_dim), dtype=x.dtype)
    for (i, grad_index) in enumerate(nonzero_grad_indices):
        all_grad[i, grad_index] = grad[grad_index]
    return all_grad
//  [python]
def before_update(self, strategy, **kwargs):
        background_num, baseball_num, bus_num, camera_num, cosplay_num, dress_num, hockey_num, laptop_num, racing_num, soccer_num, sweater_num = torch.unique(strategy.mb_y, return_counts=True)[1]
        # The compuation of beta is illustrated in our paper
        # which should be modified according to the downstream task
        beta = (baseball_num + soccer_num + hockey_num + laptop_num + camera_num) / (background_num + cosplay_num + dress_num + racing_num + sweater_num + bus_num) 
        lamda = strategy.i_batch / len(strategy.dataloader) + 1
        alpha = 1.0 * 1 ** lamda
        # x_mean = torch.mean(strategy.mb_x, 0, True)
        if strategy.experience_id != 0:
            for n, w in strategy.model.named_parameters():
                if n == "module.weight":
                    # self.pro_weight(self.Pl, x_mean, w, alpha=alpha_array, stride=2, cnn=False)
                    r = torch.mean(strategy.mb_x, 0, True)
                    k = torch.mm(self.Pl, torch.t(r))
                    # self.Pl = torch.sub(self.Pl, torch.mm(k, torch.t(k)) / (alpha + torch.mm(k, r)))
                    self.Pl = torch.sub(self.Pl, torch.mm(k, torch.t(k)) / (alpha + torch.mm(k, r)))

                    pnorm2 = torch.norm(self.Pl.data,p='fro')
                    self.Pl.data = self.Pl.data / pnorm2

                    mq = torch.mm(self.Pl, torch.linalg.solve(torch.mm(torch.t(self.Pl), self.Pl), torch.t(self.Pl)))
                    e = torch.autograd.Variable(torch.eye(len(mq)).type(torch.cuda.FloatTensor))
                    self.Ql.data = (e - mq).data                                    

                    qnorm2 = torch.norm(self.Ql.data, p='fro')
                    self.Ql.data = self.Ql.data / qnorm2
                    pwq = self.Pl + beta*self.Ql

                    w.grad.data = torch.mm(w.grad.data, torch.t(pwq.data))
//  [python]
def before_update(self, strategy, **kwargs):
        background_num, baseball_num, bus_num, camera_num, cosplay_num, dress_num, hockey_num, laptop_num, racing_num, soccer_num, sweater_num = torch.unique(strategy.mb_y, return_counts=True)[1]
        # The compuation of beta is illustrated in our paper
        # which should be modified according to the downstream task
        beta = (baseball_num + soccer_num + hockey_num + laptop_num + camera_num) / (background_num + cosplay_num + dress_num + racing_num + sweater_num + bus_num) 
        lamda = strategy.i_batch / len(strategy.dataloader) + 1
        alpha = 1.0 * 1 ** lamda
        # x_mean = torch.mean(strategy.mb_x, 0, True)
        if strategy.experience_id != 0:
            for n, w in strategy.model.named_parameters():
                if n == "module.weight":
                    # self.pro_weight(self.Pl, x_mean, w, alpha=alpha_array, stride=2, cnn=False)
                    r = torch.mean(strategy.mb_x, 0, True)
                    k = torch.mm(self.Pl, torch.t(r))
                    # self.Pl = torch.sub(self.Pl, torch.mm(k, torch.t(k)) / (alpha + torch.mm(k, r)))
                    self.Pl = torch.sub(self.Pl, torch.mm(k, torch.t(k)) / (alpha + torch.mm(k, r)))

                    pnorm2 = torch.linalg.norm(self.Pl.data,ord='fro')
                    self.Pl.data = self.Pl.data / pnorm2

                    mq = torch.mm(self.Pl, torch.linalg.solve(torch.mm(torch.t(self.Pl), self.Pl), torch.t(self.Pl)))
                    e = torch.autograd.Variable(torch.eye(len(mq)).type(torch.cuda.FloatTensor))
                    self.Ql.data = (e - mq).data                                    

                    qnorm2 = torch.linalg.norm(self.Ql.data, ord='fro')
                    self.Ql.data = self.Ql.data / qnorm2
                    pwq = self.Pl + beta*self.Ql

                    w.grad.data = torch.mm(w.grad.data, torch.t(pwq.data))
//  [python]
def calc_cosine_sim(self, dvecs, centroids):
        """
        Make the cosine similarity matrix with dims (N,M,N)
        """
        cos_sim_matrix = []
        for spkr_idx, speaker in enumerate(dvecs):
            cs_row = []
            for utt_idx, utterance in enumerate(speaker):
                new_centroids = self.calc_new_centroids(dvecs, centroids, spkr_idx, utt_idx)
                # vector based cosine similarity for speed
                cs_row.append(
                    torch.clamp(
                        torch.mm(
                            utterance.unsqueeze(1).transpose(0, 1),
                            new_centroids.transpose(0, 1),
                        )
                        / (torch.norm(utterance) * torch.norm(new_centroids, dim=1)),
                        1e-6,
                    )
                )
            cs_row = torch.cat(cs_row, dim=0)
            cos_sim_matrix.append(cs_row)
        return torch.stack(cos_sim_matrix)
//  [python]
def calc_cosine_sim(self, dvecs, centroids):
        """
        Make the cosine similarity matrix with dims (N,M,N)
        """
        cos_sim_matrix = []
        for spkr_idx, speaker in enumerate(dvecs):
            cs_row = []
            for utt_idx, utterance in enumerate(speaker):
                new_centroids = self.calc_new_centroids(dvecs, centroids, spkr_idx, utt_idx)
                # vector based cosine similarity for speed
                cs_row.append(
                    torch.clamp(
                        torch.mm(
                            utterance.unsqueeze(1).transpose(0, 1),
                            new_centroids.transpose(0, 1),
                        )
                        / (torch.linalg.norm(utterance) * torch.linalg.norm(new_centroids, dim=1)),
                        1e-6,
                    )
                )
            cs_row = torch.cat(cs_row, dim=0)
            cos_sim_matrix.append(cs_row)
        return torch.stack(cos_sim_matrix)
//  [python]
def weighted_8points(x_in, logits):
    # x_in: batch * 1 * N * 4
    x_shp = x_in.shape
    # Turn into weights for each sample
    weights = torch.relu(torch.tanh(logits))
    x_in = x_in.squeeze(1)
    
    # Make input data (num_img_pair x num_corr x 4)
    xx = torch.reshape(x_in, (x_shp[0], x_shp[2], 4)).permute(0, 2, 1)

    # Create the matrix to be used for the eight-point algorithm
    X = torch.stack([
        xx[:, 2] * xx[:, 0], xx[:, 2] * xx[:, 1], xx[:, 2],
        xx[:, 3] * xx[:, 0], xx[:, 3] * xx[:, 1], xx[:, 3],
        xx[:, 0], xx[:, 1], torch.ones_like(xx[:, 0])
    ], dim=1).permute(0, 2, 1)
    wX = torch.reshape(weights, (x_shp[0], x_shp[2], 1)) * X
    XwX = torch.matmul(X.permute(0, 2, 1), wX)
    

    # Recover essential matrix from self-adjoing eigen
    v = batch_symeig(XwX)
    e_hat = torch.reshape(v[:, :, 0], (x_shp[0], 9))

    # Make unit norm just in case
    e_hat = e_hat / torch.norm(e_hat, dim=1, keepdim=True)
    return e_hat
//  [python]
def weighted_8points(x_in, logits):
    # x_in: batch * 1 * N * 4
    x_shp = x_in.shape
    # Turn into weights for each sample
    weights = torch.relu(torch.tanh(logits))
    x_in = x_in.squeeze(1)
    
    # Make input data (num_img_pair x num_corr x 4)
    xx = torch.reshape(x_in, (x_shp[0], x_shp[2], 4)).permute(0, 2, 1)

    # Create the matrix to be used for the eight-point algorithm
    X = torch.stack([
        xx[:, 2] * xx[:, 0], xx[:, 2] * xx[:, 1], xx[:, 2],
        xx[:, 3] * xx[:, 0], xx[:, 3] * xx[:, 1], xx[:, 3],
        xx[:, 0], xx[:, 1], torch.ones_like(xx[:, 0])
    ], dim=1).permute(0, 2, 1)
    wX = torch.reshape(weights, (x_shp[0], x_shp[2], 1)) * X
    XwX = torch.matmul(X.permute(0, 2, 1), wX)
    

    # Recover essential matrix from self-adjoing eigen
    v = batch_symeig(XwX)
    e_hat = torch.reshape(v[:, :, 0], (x_shp[0], 9))

    # Make unit norm just in case
    e_hat = e_hat / torch.linalg.norm(e_hat, dim=1, keepdim=True)
    return e_hat
//  [python]
def calculate_jhmdb_PCK(predicted, gt, bbox, imgshape, thresh):

    # 0: neck    1:belly   2: face
    # 3: right shoulder  4: left shoulder
    # 5: right hip       6: left hip
    # 7: right elbow     8: left elbow
    # 9: right knee      10: left knee
    # 11: right wrist    12: left wrist
    # 13: right ankle    14: left ankle

    orderJHMDB = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14]
    N = predicted.shape[0]

    HitPoint = torch.zeros((len(orderJHMDB))).to(gt.device)
    Point_to_use = torch.ones((len(orderJHMDB))).to(gt.device)

    
    test_gt = (gt.reshape(-1, len(orderJHMDB), 2)+0.5) * imgshape[0, [1,0]]
    test_out = (predicted.reshape(-1, len(orderJHMDB), 2)+0.5) * imgshape[0, [1,0]]

    seqError = torch.zeros(N, len(orderJHMDB)).to(gt.device)
    seqThresh = torch.zeros(N, len(orderJHMDB)).to(gt.device)

    bodysize = torch.max(bbox[:, 2] - bbox[:, 0], bbox[:, 3] - bbox[:, 1])

    for frame in range(0, N):
        gt_frame = test_gt[frame, :]
        pred_frame = test_out[frame, :]

        error_dis = torch.norm(gt_frame - pred_frame,
                                   p=2,
                                   dim=1,
                                   keepdim=False)

        seqError[frame] = error_dis
        seqThresh[frame] = (bodysize[frame] * thresh) * torch.ones(
                len(orderJHMDB)).to(gt.device)

    less_than_thresh = seqError <= seqThresh
    HitPoint = torch.sum(less_than_thresh, dim=0)

    finalPCK = torch.divide(HitPoint, N)

    return finalPCK
//  [python]
def calculate_jhmdb_PCK(predicted, gt, bbox, imgshape, thresh):

    # 0: neck    1:belly   2: face
    # 3: right shoulder  4: left shoulder
    # 5: right hip       6: left hip
    # 7: right elbow     8: left elbow
    # 9: right knee      10: left knee
    # 11: right wrist    12: left wrist
    # 13: right ankle    14: left ankle

    orderJHMDB = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14]
    N = predicted.shape[0]

    HitPoint = torch.zeros((len(orderJHMDB))).to(gt.device)
    Point_to_use = torch.ones((len(orderJHMDB))).to(gt.device)

    
    test_gt = (gt.reshape(-1, len(orderJHMDB), 2)+0.5) * imgshape[0, [1,0]]
    test_out = (predicted.reshape(-1, len(orderJHMDB), 2)+0.5) * imgshape[0, [1,0]]

    seqError = torch.zeros(N, len(orderJHMDB)).to(gt.device)
    seqThresh = torch.zeros(N, len(orderJHMDB)).to(gt.device)

    bodysize = torch.max(bbox[:, 2] - bbox[:, 0], bbox[:, 3] - bbox[:, 1])

    for frame in range(0, N):
        gt_frame = test_gt[frame, :]
        pred_frame = test_out[frame, :]

        error_dis = torch.linalg.norm(gt_frame - pred_frame,
                                   ord=2,
                                   dim=1,
                                   keepdim=False)

        seqError[frame] = error_dis
        seqThresh[frame] = (bodysize[frame] * thresh) * torch.ones(
                len(orderJHMDB)).to(gt.device)

    less_than_thresh = seqError <= seqThresh
    HitPoint = torch.sum(less_than_thresh, dim=0)

    finalPCK = torch.divide(HitPoint, N)

    return finalPCK
//  [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) # N on x-axis
    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) # N on x-axis
    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 coeff2polystr(polycoeff, polyorder, A_10x10, PI_funcs_inout):
    """
    polycoeff (=Params): 
        (batch, num_params, num_functions) = (batch, 35, 32)
    polyorders: 
        (num_params, 3) = (35, 3) - degree of x,y,z for each term in polynomial.
        Since we use 4th-polynomials, the number of parameters is determined by (4+1)(4+1+1)(4+1+2)/6 = 35
    """

    batchsurfs = ["" for i in range(polycoeff.shape[0])]
    batchind = ["" for i in range(polycoeff.shape[0])]
    
    for batch in range(polycoeff.shape[0]): #batch
        surface = ""
        c = 0
        list_valid = []
        for subf in range(polycoeff.shape[2]): #100 = 25*4
            f_tmp = ""
            A_10x10[batch,:,:,subf] = A_10x10[batch,:,:,subf]/torch.norm(A_10x10[batch,:,:,subf])
            if torch.min(torch.eig(A_10x10[batch,:,:,subf])[0][:,0])<-0.0 and torch.max(torch.eig(A_10x10[batch,:,:,subf])[0][:,0])>0 and torch.sum(torch.abs(torch.prod(torch.eig(A_10x10[batch,:,:,subf])[0][:,1])))==0 and torch.min(PI_funcs_inout[batch,:,subf])<0:        
                list_valid.append(subf)

                for term in range(polycoeff.shape[1]): #35
                    f_tmp = f_tmp + '+('+str(polycoeff[batch, term, subf])+')*(x^'+str(int(polyorder[term,0]))+')*(y^'+str(int(polyorder[term,1]))+')*(z^'+str(int(polyorder[term,2]))+')'
                if c==0:
                   surface = surface + f_tmp[1:]+'\n'
                else:
                    surface = 'min(' + surface + ',' + f_tmp[1:]+')'+'\n'
                c+=1
        valids = ""
        for sub_valid in list_valid:
            valids = valids + str(sub_valid) + ' '
        batchind[batch] = batchind[batch]  + valids #+ '-' + str(R[batch,0]) 
        batchsurfs[batch] = batchsurfs[batch]  + surface #+ '-' + str(R[batch,0]) 

    return batchsurfs, batchind
//  [python]
def coeff2polystr(polycoeff, polyorder, A_10x10, PI_funcs_inout):
    """
    polycoeff (=Params): 
        (batch, num_params, num_functions) = (batch, 35, 32)
    polyorders: 
        (num_params, 3) = (35, 3) - degree of x,y,z for each term in polynomial.
        Since we use 4th-polynomials, the number of parameters is determined by (4+1)(4+1+1)(4+1+2)/6 = 35
    """

    batchsurfs = ["" for i in range(polycoeff.shape[0])]
    batchind = ["" for i in range(polycoeff.shape[0])]
    
    for batch in range(polycoeff.shape[0]): #batch
        surface = ""
        c = 0
        list_valid = []
        for subf in range(polycoeff.shape[2]): #100 = 25*4
            f_tmp = ""
            A_10x10[batch,:,:,subf] = A_10x10[batch,:,:,subf]/torch.linalg.norm(A_10x10[batch,:,:,subf])
            if torch.min(torch.eig(A_10x10[batch,:,:,subf])[0][:,0])<-0.0 and torch.max(torch.eig(A_10x10[batch,:,:,subf])[0][:,0])>0 and torch.sum(torch.abs(torch.prod(torch.eig(A_10x10[batch,:,:,subf])[0][:,1])))==0 and torch.min(PI_funcs_inout[batch,:,subf])<0:        
                list_valid.append(subf)

                for term in range(polycoeff.shape[1]): #35
                    f_tmp = f_tmp + '+('+str(polycoeff[batch, term, subf])+')*(x^'+str(int(polyorder[term,0]))+')*(y^'+str(int(polyorder[term,1]))+')*(z^'+str(int(polyorder[term,2]))+')'
                if c==0:
                   surface = surface + f_tmp[1:]+'\n'
                else:
                    surface = 'min(' + surface + ',' + f_tmp[1:]+')'+'\n'
                c+=1
        valids = ""
        for sub_valid in list_valid:
            valids = valids + str(sub_valid) + ' '
        batchind[batch] = batchind[batch]  + valids #+ '-' + str(R[batch,0]) 
        batchsurfs[batch] = batchsurfs[batch]  + surface #+ '-' + str(R[batch,0]) 

    return batchsurfs, batchind
//  [python]
def calc_str_loss(pred, true, 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)
    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)
    loss = -(torch.nn.functional.relu(1.0-difference/A)).mean(dim=(1,2,3)) # (I)
    
    # 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()
    #tot_loss = tot_loss + loss[-1] # emphasize last results more!
    return tot_loss, loss.detach()
//  [python]
def calc_str_loss(pred, true, 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)
    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)
    loss = -(torch.nn.functional.relu(1.0-difference/A)).mean(dim=(1,2,3)) # (I)
    
    # 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()
    #tot_loss = tot_loss + loss[-1] # emphasize last results more!
    return tot_loss, loss.detach()
//  [python]
def forward_train(self, x: Tensor, gt_label: Tensor, **kwargs) -> Dict:
        """Forward training data."""
        x_norm = torch.norm(x, p=2, dim=1).unsqueeze(1).expand_as(x)
        x_normalized = x.div(x_norm + self.eps)
        cls_score = self.temperature * self.fc(x_normalized)
        losses = self.loss(cls_score, gt_label)
        return losses
//  [python]
def forward_train(self, x: Tensor, gt_label: Tensor, **kwargs) -> Dict:
        """Forward training data."""
        x_norm = torch.linalg.norm(x, ord=2, dim=1).unsqueeze(1).expand_as(x)
        x_normalized = x.div(x_norm + self.eps)
        cls_score = self.temperature * self.fc(x_normalized)
        losses = self.loss(cls_score, gt_label)
        return losses
//  [python]
def forward_query(self, x: Tensor, **kwargs) -> List:
        """Forward query data in meta testing."""
        x_norm = torch.norm(x, p=2, dim=1).unsqueeze(1).expand_as(x)
        x_normalized = x.div(x_norm + self.eps)
        cls_score = self.temperature * self.fc(x_normalized)
        pred = F.softmax(cls_score, dim=1)
        pred = list(pred.detach().cpu().numpy())
        return pred
//  [python]
def forward_query(self, x: Tensor, **kwargs) -> List:
        """Forward query data in meta testing."""
        x_norm = torch.linalg.norm(x, ord=2, dim=1).unsqueeze(1).expand_as(x)
        x_normalized = x.div(x_norm + self.eps)
        cls_score = self.temperature * self.fc(x_normalized)
        pred = F.softmax(cls_score, dim=1)
        pred = list(pred.detach().cpu().numpy())
        return pred
//  [python]
def infer(self, data, dictionary):
        """
        Infer coefficients for each image in data using dict elements dictionary using Matching Pursuit (MP)

        Parameters
        ----------
        data : array-like (batch_size, n_features)
            data to be used in sparse coding
        dictionary : array-like, (n_features, n_basis)
            dictionary to be used to get the coefficients

        Returns
        -------
        coefficients : array-like (batch_size, n_basis)
        """
        # Get input characteristics
        batch_size, n_features = data.shape
        n_features, n_basis = dictionary.shape
        device = dictionary.device

        # Define signal sparsity
        K = np.ceil(self.sparsity*n_basis).astype(int)

        # Get dictionary norms in case atoms are not normalized
        dictionary_norms = torch.norm(dictionary, p=2, dim=0, keepdim=True)

        # Initialize coefficients for the whole batch
        coefficients = torch.zeros(
            batch_size, n_basis, requires_grad=False, device=device)

        residual = data.clone()  # [batch_size, n_features]

        for _ in range(K):
            # Select which (coefficient, basis function) pair to update using the inner product.
            candidate_coefs = residual @ dictionary  # [batch_size, n_basis]
            top_coef_idxs = torch.argmax(torch.abs(candidate_coefs) / dictionary_norms, dim=1)  # [batch_size]

            # Update the coefficient.
            top_coefs = candidate_coefs[torch.arange(batch_size), top_coef_idxs]  # [batch_size]
            coefficients[torch.arange(batch_size), top_coef_idxs] = top_coefs

            # Explain away/subtract the chosen coefficient and corresponding basis from the residual.
            top_coef_bases = dictionary[..., top_coef_idxs]  # [n_features, batch_size]
            residual = residual - top_coefs.reshape(batch_size, 1) * top_coef_bases.T  # [batch_size, n_features]

        return coefficients.detach()
//  [python]
def infer(self, data, dictionary):
        """
        Infer coefficients for each image in data using dict elements dictionary using Matching Pursuit (MP)

        Parameters
        ----------
        data : array-like (batch_size, n_features)
            data to be used in sparse coding
        dictionary : array-like, (n_features, n_basis)
            dictionary to be used to get the coefficients

        Returns
        -------
        coefficients : array-like (batch_size, n_basis)
        """
        # Get input characteristics
        batch_size, n_features = data.shape
        n_features, n_basis = dictionary.shape
        device = dictionary.device

        # Define signal sparsity
        K = np.ceil(self.sparsity*n_basis).astype(int)

        # Get dictionary norms in case atoms are not normalized
        dictionary_norms = torch.linalg.norm(dictionary, p=2, dim=0, keepdim=True)

        # Initialize coefficients for the whole batch
        coefficients = torch.zeros(
            batch_size, n_basis, requires_grad=False, device=device)

        residual = data.clone()  # [batch_size, n_features]

        for _ in range(K):
            # Select which (coefficient, basis function) pair to update using the inner product.
            candidate_coefs = residual @ dictionary  # [batch_size, n_basis]
            top_coef_idxs = torch.argmax(torch.abs(candidate_coefs) / dictionary_norms, dim=1)  # [batch_size]

            # Update the coefficient.
            top_coefs = candidate_coefs[torch.arange(batch_size), top_coef_idxs]  # [batch_size]
            coefficients[torch.arange(batch_size), top_coef_idxs] = top_coefs

            # Explain away/subtract the chosen coefficient and corresponding basis from the residual.
            top_coef_bases = dictionary[..., top_coef_idxs]  # [n_features, batch_size]
            residual = residual - top_coefs.reshape(batch_size, 1) * top_coef_bases.T  # [batch_size, n_features]

        return coefficients.detach()
//  [python]
def intersect(self, rays, distance):
        distance = distance.view(distance.shape[0], -1)

        # Base distance
        rays_o, rays_d = rays[..., :3], rays[..., 3:6]
        base_pos = self.unify_fn(rays)
        diff = base_pos - rays_o

        # Add distance offset
        distance = distance + (
            torch.sign(dot(rays_d, diff)) * torch.norm(diff, dim=-1)
        ).unsqueeze(1)

        return distance
//  [python]
def intersect(self, rays, distance):
        distance = distance.view(distance.shape[0], -1)

        # Base distance
        rays_o, rays_d = rays[..., :3], rays[..., 3:6]
        base_pos = self.unify_fn(rays)
        diff = base_pos - rays_o

        # Add distance offset
        distance = distance + (
            torch.sign(dot(rays_d, diff)) * torch.linalg.norm(diff, dim=-1)
        ).unsqueeze(1)

        return distance
//  [python]
print('Are these both log(det(A))?', bool(logdet1 - logdet2 < 1e-8))

logdet1.backward()
logdet2.backward()

inv_mat = mat.inverse()
print('Does Way 1 yield A^{-1}?', bool(torch.norm(mat.grad - inv_mat) < 1e-8))
//  [python]
print('Are these both log(det(A))?', bool(logdet1 - logdet2 < 1e-8))

logdet1.backward()
logdet2.backward()

inv_mat = mat.inverse()
print('Does Way 1 yield A^{-1}?', bool(torch.linalg.norm(mat.grad - inv_mat) < 1e-8))
//  [python]
def infer(self, data, dictionary):
        """
        Infer coefficients for each image in data using dict elements dictionary using Orthogonal Matching Pursuit (OMP)

        Parameters
        ----------
        data : array-like (batch_size, n_features)
            data to be used in sparse coding
        dictionary : array-like, (n_features, n_basis)
            dictionary to be used to get the coefficients

        Returns
        -------
        coefficients : array-like (batch_size, n_basis)
        """
        # Get input characteristics
        batch_size, n_features = data.shape
        n_features, n_basis = dictionary.shape
        device = dictionary.device

        # Define signal sparsity
        K = np.ceil(self.sparsity*n_basis).astype(int)

        # Get dictionary norms in case atoms are not normalized
        dictionary_norms = torch.norm(dictionary, p=2, dim=0, keepdim=True)

        # Initialize coefficients for the whole batch
        coefficients = torch.zeros(
            batch_size, n_basis, requires_grad=False, device=device)

        residual = data.clone()  # [batch_size, n_features]

        # The basis functions that are used to infer the coefficients will be updated each time.
        used_basis_fns = torch.zeros((batch_size, n_basis), dtype=torch.bool)

        for t in range(K):
            # Select which (coefficient, basis function) pair to update using the inner product.
            candidate_coefs = residual @ dictionary  # [batch_size, n_basis]
            top_coef_idxs = torch.argmax(torch.abs(candidate_coefs) / dictionary_norms, dim=1)  # [batch_size]
            used_basis_fns[torch.arange(batch_size), top_coef_idxs] = True

            # Update the coefficients
            used_dictionary = dictionary[..., used_basis_fns.nonzero()[:, 1]].reshape((n_features, batch_size, t + 1))

            (used_coefficients, _, _, _) = torch.linalg.lstsq(
                used_dictionary.permute((1, 0, 2)),  # [batch_size, n_features, t + 1]
                data.reshape(batch_size, n_features, 1),
            )  # [batch_size, t + 1, 1]
            coefficients[used_basis_fns] = used_coefficients.reshape(-1)

            # Update the residual.
            residual = data.clone() - coefficients @ dictionary.T  # [batch_size, n_features]

        return coefficients.detach()
//  [python]
def infer(self, data, dictionary):
        """
        Infer coefficients for each image in data using dict elements dictionary using Orthogonal Matching Pursuit (OMP)

        Parameters
        ----------
        data : array-like (batch_size, n_features)
            data to be used in sparse coding
        dictionary : array-like, (n_features, n_basis)
            dictionary to be used to get the coefficients

        Returns
        -------
        coefficients : array-like (batch_size, n_basis)
        """
        # Get input characteristics
        batch_size, n_features = data.shape
        n_features, n_basis = dictionary.shape
        device = dictionary.device

        # Define signal sparsity
        K = np.ceil(self.sparsity*n_basis).astype(int)

        # Get dictionary norms in case atoms are not normalized
        dictionary_norms = torch.linalg.norm(dictionary, p=2, dim=0, keepdim=True)

        # Initialize coefficients for the whole batch
        coefficients = torch.zeros(
            batch_size, n_basis, requires_grad=False, device=device)

        residual = data.clone()  # [batch_size, n_features]

        # The basis functions that are used to infer the coefficients will be updated each time.
        used_basis_fns = torch.zeros((batch_size, n_basis), dtype=torch.bool)

        for t in range(K):
            # Select which (coefficient, basis function) pair to update using the inner product.
            candidate_coefs = residual @ dictionary  # [batch_size, n_basis]
            top_coef_idxs = torch.argmax(torch.abs(candidate_coefs) / dictionary_norms, dim=1)  # [batch_size]
            used_basis_fns[torch.arange(batch_size), top_coef_idxs] = True

            # Update the coefficients
            used_dictionary = dictionary[..., used_basis_fns.nonzero()[:, 1]].reshape((n_features, batch_size, t + 1))

            (used_coefficients, _, _, _) = torch.linalg.lstsq(
                used_dictionary.permute((1, 0, 2)),  # [batch_size, n_features, t + 1]
                data.reshape(batch_size, n_features, 1),
            )  # [batch_size, t + 1, 1]
            coefficients[used_basis_fns] = used_coefficients.reshape(-1)

            # Update the residual.
            residual = data.clone() - coefficients @ dictionary.T  # [batch_size, n_features]

        return coefficients.detach()
//  [python]
def predFlowCoarse(corrKernel21, NetFlowCoarse, grid, up8X = True) : 
    
    flowCoarse = NetFlowCoarse(corrKernel21, up8X) ## output is with dimension B, 2, W, H
    b, _, w, h = flowCoarse.size()
    flowGrad = flowCoarse.narrow(2, 1, w-1).narrow(3, 1, h-1) - flowCoarse.narrow(2, 0, w-1).narrow(3, 0, h-1)
    flowGrad = torch.norm(flowGrad, dim=1, keepdim=True)
    flowCoarse = flowCoarse.permute(0, 2, 3, 1)
    flowCoarse = torch.clamp(flowCoarse + grid, min=-1, max=1)
    
    return flowGrad, flowCoarse
//  [python]
def predFlowCoarse(corrKernel21, NetFlowCoarse, grid, up8X = True) : 
    
    flowCoarse = NetFlowCoarse(corrKernel21, up8X) ## output is with dimension B, 2, W, H
    b, _, w, h = flowCoarse.size()
    flowGrad = flowCoarse.narrow(2, 1, w-1).narrow(3, 1, h-1) - flowCoarse.narrow(2, 0, w-1).narrow(3, 0, h-1)
    flowGrad = torch.linalg.norm(flowGrad, dim=1, keepdim=True)
    flowCoarse = flowCoarse.permute(0, 2, 3, 1)
    flowCoarse = torch.clamp(flowCoarse + grid, min=-1, max=1)
    
    return flowGrad, flowCoarse
//  [python]
def intersect(self, rays, z_vals):
        z_vals = z_vals.view(z_vals.shape[0], self.z_channels, 8)
        origins = z_vals[..., :3]
        resize = z_vals[..., 3:6]
        raw_offsets = z_vals[..., -2]
        radii = z_vals[..., -1]

        # Transform the space
        rays_o = (rays[..., None, 0:3] - origins) * resize
        rays_d = rays[..., None, 3:6] * resize
        rays = torch.cat(
            [
                rays_o,
                torch.nn.functional.normalize(rays_d, p=2.0, dim=-1),
            ],
            -1
        )

        # Calculate intersection distances (in transformed space)
        t = intersect_sphere(
            rays,
            torch.zeros_like(origins),
            radii,
        )

        # Recycle samples for not-hit spheres
        min_radius = min_sphere_radius(rays, torch.zeros_like(rays[..., :3]))

        base_pos = self.unify_fn(rays)
        diff = (base_pos - rays[..., :3])
        base_distance = torch.sign(dot(rays[..., 3:6], diff)) * torch.norm(diff, dim=-1)

        t = torch.where(
            torch.abs(radii) < min_radius + 4 * self.z_scale,
            raw_offsets + base_distance,
            t
        )

        # Transform distances
        return t / (torch.norm(rays_d, dim=-1) + 1e-5)
//  [python]
def intersect(self, rays, z_vals):
        z_vals = z_vals.view(z_vals.shape[0], self.z_channels, 8)
        origins = z_vals[..., :3]
        resize = z_vals[..., 3:6]
        raw_offsets = z_vals[..., -2]
        radii = z_vals[..., -1]

        # Transform the space
        rays_o = (rays[..., None, 0:3] - origins) * resize
        rays_d = rays[..., None, 3:6] * resize
        rays = torch.cat(
            [
                rays_o,
                torch.nn.functional.normalize(rays_d, p=2.0, dim=-1),
            ],
            -1
        )

        # Calculate intersection distances (in transformed space)
        t = intersect_sphere(
            rays,
            torch.zeros_like(origins),
            radii,
        )

        # Recycle samples for not-hit spheres
        min_radius = min_sphere_radius(rays, torch.zeros_like(rays[..., :3]))

        base_pos = self.unify_fn(rays)
        diff = (base_pos - rays[..., :3])
        base_distance = torch.sign(dot(rays[..., 3:6], diff)) * torch.linalg.norm(diff, dim=-1)

        t = torch.where(
            torch.abs(radii) < min_radius + 4 * self.z_scale,
            raw_offsets + base_distance,
            t
        )

        # Transform distances
        return t / (torch.linalg.norm(rays_d, dim=-1) + 1e-5)
//  [python]
def intersect(self, rays, z_vals):
        z_vals = z_vals.view(z_vals.shape[0], self.z_channels, 8)
        origins = z_vals[..., :3]
        resize = z_vals[..., 3:6]
        raw_offsets = z_vals[..., -2]
        radii = z_vals[..., -1]

        # Transform the space
        rays_o = (rays[..., None, 0:3] - origins) * resize
        rays_d = rays[..., None, 3:6] * resize
        rays = torch.cat(
            [
                rays_o,
                torch.nn.functional.normalize(rays_d, p=2.0, dim=-1),
            ],
            -1
        )

        # Calculate intersection distances (in transformed space)
        t = intersect_cylinder(
            rays,
            torch.zeros_like(origins),
            radii,
        )

        # Recycle samples for not-hit cylinders
        min_radius = min_cylinder_radius(rays, torch.zeros_like(rays[..., :3]))

        base_pos = self.unify_fn(rays)
        rays_o_cyl = torch.cat([rays[..., 0:1], torch.zeros_like(rays[..., 1:2]), rays[..., 2:3]], -1)
        rays_d_cyl = torch.cat([rays[..., 3:4], torch.zeros_like(rays[..., 4:5]), rays[..., 5:6]], -1)
        diff = (base_pos - rays_o_cyl)

        base_distance = torch.sign(dot(rays_d_cyl, diff)) * torch.norm(diff, dim=-1) / torch.norm(rays_d_cyl, dim=-1)

        t = torch.where(
            torch.abs(radii) < min_radius + 4 * self.z_scale,
            raw_offsets + base_distance,
            t
        )

        # Transform distances
        return t / (torch.norm(rays_d, dim=-1) + 1e-5)
//  [python]
def intersect(self, rays, z_vals):
        z_vals = z_vals.view(z_vals.shape[0], self.z_channels, 8)
        origins = z_vals[..., :3]
        resize = z_vals[..., 3:6]
        raw_offsets = z_vals[..., -2]
        radii = z_vals[..., -1]

        # Transform the space
        rays_o = (rays[..., None, 0:3] - origins) * resize
        rays_d = rays[..., None, 3:6] * resize
        rays = torch.cat(
            [
                rays_o,
                torch.nn.functional.normalize(rays_d, p=2.0, dim=-1),
            ],
            -1
        )

        # Calculate intersection distances (in transformed space)
        t = intersect_cylinder(
            rays,
            torch.zeros_like(origins),
            radii,
        )

        # Recycle samples for not-hit cylinders
        min_radius = min_cylinder_radius(rays, torch.zeros_like(rays[..., :3]))

        base_pos = self.unify_fn(rays)
        rays_o_cyl = torch.cat([rays[..., 0:1], torch.zeros_like(rays[..., 1:2]), rays[..., 2:3]], -1)
        rays_d_cyl = torch.cat([rays[..., 3:4], torch.zeros_like(rays[..., 4:5]), rays[..., 5:6]], -1)
        diff = (base_pos - rays_o_cyl)

        base_distance = torch.sign(dot(rays_d_cyl, diff)) * torch.linalg.norm(diff, dim=-1) / torch.linalg.norm(rays_d_cyl, dim=-1)

        t = torch.where(
            torch.abs(radii) < min_radius + 4 * self.z_scale,
            raw_offsets + base_distance,
            t
        )

        # Transform distances
        return t / (torch.linalg.norm(rays_d, dim=-1) + 1e-5)
//  [python]
def log_det(y_true, y_pred, num_model):
    num_classes = 10
    bool_R_y_true = torch.ne(torch.ones(y_true.shape) - y_true, 0)  # batch_size X (num_class X num_models), 2-D
    mask_non_y_pred = y_pred[bool_R_y_true]  # batch_size X (num_class-1) X num_models, 1-D
    mask_non_y_pred = torch.reshape(mask_non_y_pred, [-1, num_model, num_classes - 1])  # batch_size X num_model X (num_class-1), 3-D
    mask_non_y_pred = mask_non_y_pred / torch.norm(mask_non_y_pred, dim=2, keepdim=True)  # batch_size X num_model X (num_class-1), 3-D
    matrix = torch.bmm(mask_non_y_pred, torch.transpose(mask_non_y_pred, 1, 2))  # batch_size X num_model X num_model, 3-D
    all_log_det = torch.log(torch.linalg.det(matrix + det_offset * torch.unsqueeze(torch.eye(num_model), 0)))  # batch_size X 1, 1-D

    # bool_R_y_true = tf.not_equal(tf.ones_like(y_true) - y_true, zero) # batch_size X (num_class X num_models), 2-D
    # mask_non_y_pred = tf.boolean_mask(y_pred, bool_R_y_true) # batch_size X (num_class-1) X num_models, 1-D
    # mask_non_y_pred = tf.reshape(mask_non_y_pred, [-1, num_model, num_classes-1]) # batch_size X num_model X (num_class-1), 3-D
    # mask_non_y_pred = mask_non_y_pred / tf.norm(mask_non_y_pred, axis=2, keepdims=True) # batch_size X num_model X (num_class-1), 3-D
    # matrix = tf.matmul(mask_non_y_pred, tf.transpose(mask_non_y_pred, perm=[0, 2, 1])) # batch_size X num_model X num_model, 3-D
    # all_log_det = tf.linalg.logdet(matrix+det_offset*tf.expand_dims(tf.eye(num_model),0)) # batch_size X 1, 1-D
    return all_log_det
//  [python]
def log_det(y_true, y_pred, num_model):
    num_classes = 10
    bool_R_y_true = torch.ne(torch.ones(y_true.shape) - y_true, 0)  # batch_size X (num_class X num_models), 2-D
    mask_non_y_pred = y_pred[bool_R_y_true]  # batch_size X (num_class-1) X num_models, 1-D
    mask_non_y_pred = torch.reshape(mask_non_y_pred, [-1, num_model, num_classes - 1])  # batch_size X num_model X (num_class-1), 3-D
    mask_non_y_pred = mask_non_y_pred / torch.linalg.norm(mask_non_y_pred, dim=2, keepdim=True)  # batch_size X num_model X (num_class-1), 3-D
    matrix = torch.bmm(mask_non_y_pred, torch.transpose(mask_non_y_pred, 1, 2))  # batch_size X num_model X num_model, 3-D
    all_log_det = torch.log(torch.linalg.det(matrix + det_offset * torch.unsqueeze(torch.eye(num_model), 0)))  # batch_size X 1, 1-D
    return all_log_det
//  [python]
def normalize_vecs(vectors: torch.Tensor) ->torch.Tensor:
    """
    Normalize vector lengths.
    """
    return vectors / torch.norm(vectors, dim=-1, keepdim=True)
//  [python]
def normalize_vecs(vectors: torch.Tensor) ->torch.Tensor:
    """
    Normalize vector lengths.
    """
    return vectors / torch.linalg.norm(vectors, dim=-1, keepdim=True)
//  [python]
def _orthogonalize(matrix, epsilon=0):
    """
    Applies Gram-Schmidt procedure to orthogonalize a given 2D tensor.
    If epsilon is 0, this is equivalent to `torch.qr(matrix, out=(matrix, _))`,
    """
    # TODO Consider using Q = torch.orgqr(*torch.geqrf(A)) to compute the Q of the QR _much_ faster
    # and more reliably.
    # Works on FP32/64 or complex numbers (does not work for half precision)
    num_cols = matrix.shape[1]
    for i in range(num_cols):
        # Normalize the i'th column.
        col = matrix[:, i : i + 1]
        # If no epsilon is added here, division by zero may be caused by vanishing gradients.
        # This epsilon is not needed if the input matrix covers the gradients of at least one entire layer in the neural network.
        if epsilon == 0:
            # Note that col ** 2 can underflow/overflow if we use FP16.
            # May need to consider multiplying a scaling factor and dividing it later, or using bfloat16 instead.
            try:
                col /= torch.norm(col)
            except ZeroDivisionError:
                logging.error(
                    "The matrix to be orthogonalized has at least a column of all 0s. Please set a small value such as 1e-8 "
                    "as `orthogonalization_epsilon` in PowerSGD state."
                )
                # Recover the values from NaNs to 0s.
                col.fill_(0.0)
        else:
            col /= torch.norm(col) + epsilon
        # Project it on the rest and remove it.
        if i + 1 < num_cols:
            rest = matrix[:, i + 1 :]
            rest -= torch.sum(col * rest, dim=0) * col