jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def test_chain_matmul(self, device, dtype):
        # chain_matmul accepts a single input tensor while multi_dot does not
        t = make_tensor((2, 2), device, dtype)
        self.assertEqual(t, torch.chain_matmul(t))
        with self.assertRaisesRegex(RuntimeError, r"chain_matmul\(\): Expected one or more matrices"):
            torch.chain_matmul()

        # chain_matmul expects all tensors to be 2D whereas multi_dot allows the first and last tensors to
        # be either 1D or 2D
        with self.assertRaisesRegex(RuntimeError, r"Tensor dimension is 1, expected 2 instead"):
            torch.chain_matmul(make_tensor(1, device, dtype), make_tensor(1, device, dtype))
//  [python]
def test_chain_matmul(self, device, dtype):
        # chain_matmul accepts a single input tensor while multi_dot does not
        t = make_tensor((2, 2), device, dtype)
        self.assertEqual(t, torch.linalg.multi_dot([t]))
        with self.assertRaisesRegex(RuntimeError, r"chain_matmul\(\): Expected one or more matrices"):
            torch.linalg.multi_dot([])

        # chain_matmul expects all tensors to be 2D whereas multi_dot allows the first and last tensors to
        # be either 1D or 2D
        with self.assertRaisesRegex(RuntimeError, r"Tensor dimension is 1, expected 2 instead"):
            torch.linalg.multi_dot([make_tensor(1, device, dtype), make_tensor(1, device, dtype)])
//  [python]
def _solve_v_and_rescale(self, weight_mat, u, target_sigma):
        v = torch.chain_matmul(weight_mat.t().mm(weight_mat).pinverse(), weight_mat.t(), u.unsqueeze(1)).squeeze(1)
        return v.mul_(target_sigma / torch.dot(u, torch.mv(weight_mat, v)))
//  [python]
def _solve_v_and_rescale(self, weight_mat, u, target_sigma):
        v = torch.linalg.multi_dot([weight_mat.t().mm(weight_mat).pinverse(), weight_mat.t(), u.unsqueeze(1)]).squeeze(1)
        return v.mul_(target_sigma / torch.dot(u, torch.mv(weight_mat, v)))
//  [python]
def frobgrad(matrices):

    assert type(matrices) == list
    assert 1 <= len(matrices) <= 3
    output = [None for _ in matrices]
    if len(matrices) == 1:
        output[0] = matrices[0].clone()
    else:
        if len(matrices) == 2:
            U, VT = matrices
            UM, MVT = matrices
        else:
            U, M, VT = matrices
            UM, MVT = torch.matmul(U, M), torch.matmul(M, VT)
            output[1] = torch.chain_matmul(U.T, U, M, VT, VT.T)
        output[0] = torch.chain_matmul(U, MVT, MVT.T)
        output[-1] = torch.chain_matmul(UM.T, UM, VT)
    return output
//  [python]
def frobgrad(matrices):

    assert type(matrices) == list
    assert 1 <= len(matrices) <= 3
    output = [None for _ in matrices]
    if len(matrices) == 1:
        output[0] = matrices[0].clone()
    else:
        if len(matrices) == 2:
            U, VT = matrices
            UM, MVT = matrices
        else:
            U, M, VT = matrices
            UM, MVT = torch.matmul(U, M), torch.matmul(M, VT)
            output[1] = torch.linalg.multi_dot([U.T, U, M, VT, VT.T])
        output[0] = torch.linalg.multi_dot([U, MVT, MVT.T])
        output[-1] = torch.linalg.multi_dot([UM.T, UM, VT])
    return output
//  [python]
def warp(
        pos1,
        depth1, intrinsics1, pose1, bbox1,
        depth2, intrinsics2, pose2, bbox2
):
    device = pos1.device

    Z1, pos1, ids = interpolate_depth(pos1, depth1)

    # COLMAP convention
    u1 = pos1[1, :] + bbox1[1] + .5
    v1 = pos1[0, :] + bbox1[0] + .5

    X1 = (u1 - intrinsics1[0, 2]) * (Z1 / intrinsics1[0, 0])
    Y1 = (v1 - intrinsics1[1, 2]) * (Z1 / intrinsics1[1, 1])

    XYZ1_hom = torch.cat([
        X1.view(1, -1),
        Y1.view(1, -1),
        Z1.view(1, -1),
        torch.ones(1, Z1.size(0), device=device)
    ], dim=0)
    XYZ2_hom = torch.chain_matmul(pose2, torch.inverse(pose1), XYZ1_hom)
    XYZ2 = XYZ2_hom[: -1, :] / XYZ2_hom[-1, :].view(1, -1)

    uv2_hom = torch.matmul(intrinsics2, XYZ2)
    uv2 = uv2_hom[: -1, :] / uv2_hom[-1, :].view(1, -1)

    u2 = uv2[0, :] - bbox2[1] - .5
    v2 = uv2[1, :] - bbox2[0] - .5
    uv2 = torch.cat([u2.view(1, -1),  v2.view(1, -1)], dim=0)

    annotated_depth, pos2, new_ids = interpolate_depth(uv_to_pos(uv2), depth2)

    ids = ids[new_ids]
    pos1 = pos1[:, new_ids]
    estimated_depth = XYZ2[2, new_ids]

    inlier_mask = torch.abs(estimated_depth - annotated_depth) < 0.05

    ids = ids[inlier_mask]
    if ids.size(0) == 0:
        raise EmptyTensorError

    pos2 = pos2[:, inlier_mask]
    pos1 = pos1[:, inlier_mask]

    return pos1, pos2, ids
//  [python]
def warp(
        pos1,
        depth1, intrinsics1, pose1, bbox1,
        depth2, intrinsics2, pose2, bbox2
):
    device = pos1.device

    Z1, pos1, ids = interpolate_depth(pos1, depth1)

    # COLMAP convention
    u1 = pos1[1, :] + bbox1[1] + .5
    v1 = pos1[0, :] + bbox1[0] + .5

    X1 = (u1 - intrinsics1[0, 2]) * (Z1 / intrinsics1[0, 0])
    Y1 = (v1 - intrinsics1[1, 2]) * (Z1 / intrinsics1[1, 1])

    XYZ1_hom = torch.cat([
        X1.view(1, -1),
        Y1.view(1, -1),
        Z1.view(1, -1),
        torch.ones(1, Z1.size(0), device=device)
    ], dim=0)
    XYZ2_hom = torch.linalg.multi_dot([pose2, torch.inverse(pose1), XYZ1_hom])
    XYZ2 = XYZ2_hom[: -1, :] / XYZ2_hom[-1, :].view(1, -1)

    uv2_hom = torch.matmul(intrinsics2, XYZ2)
    uv2 = uv2_hom[: -1, :] / uv2_hom[-1, :].view(1, -1)

    u2 = uv2[0, :] - bbox2[1] - .5
    v2 = uv2[1, :] - bbox2[0] - .5
    uv2 = torch.cat([u2.view(1, -1),  v2.view(1, -1)], dim=0)

    annotated_depth, pos2, new_ids = interpolate_depth(uv_to_pos(uv2), depth2)

    ids = ids[new_ids]
    pos1 = pos1[:, new_ids]
    estimated_depth = XYZ2[2, new_ids]

    inlier_mask = torch.abs(estimated_depth - annotated_depth) < 0.05

    ids = ids[inlier_mask]
    if ids.size(0) == 0:
        raise EmptyTensorError

    pos2 = pos2[:, inlier_mask]
    pos1 = pos1[:, inlier_mask]

    return pos1, pos2, ids
//  [python]
def sample_energy(ceta, mean, cov, zi,n_gmm,bs):
        # print('calculate sample energy')
        e = torch.tensor(0.0)
        cov_eps = torch.eye(mean.shape[1]) * (1e-3) # original constant: 1e-12
#         cov_eps = cov_eps.to(device)
        for k in range(n_gmm):
            miu_k = mean[k].unsqueeze(1)
            d_k = zi - miu_k

            inv_cov = torch.inverse(cov[k] + cov_eps)
            
           
            # inv_cov = inv_cov.detach().numpy()
            # with open('test_log.txt', 'a') as f:
            #     if inv_cov.shape[0] != np.linalg.matrix_rank(inv_cov):
            #         f.write('matrix does not have full rank. \n')
            #     else:
            #         f.write('matrix have full rank. \n')
            #     f.close()
            # inv_cov = torch.tensor(inv_cov)
            # inv_cov.to('cuda')
            
            e_k = torch.exp(-0.5 * torch.chain_matmul(torch.t(d_k), inv_cov, d_k))
            e_k = e_k / torch.sqrt(torch.abs(torch.det(2 * math.pi * cov[k])))
            e_k = e_k * ceta[k]
            e += e_k.squeeze()
            
        return -torch.log(e)
//  [python]
def sample_energy(ceta, mean, cov, zi,n_gmm,bs):
        # print('calculate sample energy')
        e = torch.tensor(0.0)
        cov_eps = torch.eye(mean.shape[1]) * (1e-3) # original constant: 1e-12
#         cov_eps = cov_eps.to(device)
        for k in range(n_gmm):
            miu_k = mean[k].unsqueeze(1)
            d_k = zi - miu_k

            inv_cov = torch.inverse(cov[k] + cov_eps)
            
           
            # inv_cov = inv_cov.detach().numpy()
            # with open('test_log.txt', 'a') as f:
            #     if inv_cov.shape[0] != np.linalg.matrix_rank(inv_cov):
            #         f.write('matrix does not have full rank. \n')
            #     else:
            #         f.write('matrix have full rank. \n')
            #     f.close()
            # inv_cov = torch.tensor(inv_cov)
            # inv_cov.to('cuda')
            
            e_k = torch.exp(-0.5 * torch.linalg.multi_dot([torch.t(d_k), inv_cov, d_k]))
            e_k = e_k / torch.sqrt(torch.abs(torch.det(2 * math.pi * cov[k])))
            e_k = e_k * ceta[k]
            e += e_k.squeeze()
            
        return -torch.log(e)
//  [python]
def inverse(self):
        """
            fast inverse computation making use of matrix structure
        :return:
        """
        with torch.no_grad():
            # t0 = time.time()
            self.LU_decomposition()
            # t1 = time.time()
            submat_inv = self.tri_submatrix_inverse()
            # t2 = time.time()
            reci = 1.0 / \
                   (self.last_diag - torch.chain_matmul(self.last_vec.view(1, -1), submat_inv, self.last_vec.view(-1, 1)).item())
            inv_vec = - reci * torch.matmul(submat_inv, self.last_vec)

            ans = torch.zeros(self.n, self.n)
            ans[:-1, :-1] = submat_inv + \
                            reci * torch.chain_matmul(submat_inv, self.last_vec.view(-1, 1),
                                                      self.last_vec.view(1, -1), submat_inv)
            ans[-1, :-1] = inv_vec
            ans[:-1, -1] = inv_vec
            ans[-1, -1] = reci
            # t3 = time.time()

        # print(t1 - t0, '+', t2 - t1, '+', t3 - t2)

        self.inv = ans
        return ans
//  [python]
def inverse(self):
        """
            fast inverse computation making use of matrix structure
        :return:
        """
        with torch.no_grad():
            # t0 = time.time()
            self.LU_decomposition()
            # t1 = time.time()
            submat_inv = self.tri_submatrix_inverse()
            # t2 = time.time()
            reci = 1.0 / \
                   (self.last_diag - torch.linalg.multi_dot([self.last_vec.view(1, -1), submat_inv, self.last_vec.view(-1, 1)]).item())
            inv_vec = - reci * torch.matmul(submat_inv, self.last_vec)

            ans = torch.zeros(self.n, self.n)
            ans[:-1, :-1] = submat_inv + \
                            reci * torch.linalg.multi_dot([submat_inv, self.last_vec.view(-1, 1),
                                                      self.last_vec.view(1, -1), submat_inv])
            ans[-1, :-1] = inv_vec
            ans[:-1, -1] = inv_vec
            ans[-1, -1] = reci
            # t3 = time.time()

        # print(t1 - t0, '+', t2 - t1, '+', t3 - t2)

        self.inv = ans
        return ans
//  [python]
def forward(self, W, U0, ms, d, projector='sinkhorn'):
        num_graphs = ms.shape[0]

        U = U0
        for i in range(self.max_iter):
            lastU = U
            WU = torch.mm(W, U) #/ num_graphs
            V = torch.chain_matmul(WU, U.t(), WU) #/ num_graphs ** 2

            #V_median = torch.median(torch.flatten(V, start_dim=-2), dim=-1).values
            #V_var, V_mean = torch.var_mean(torch.flatten(V, start_dim=-2), dim=-1)
            #V = V - V_mean
            #V = V / torch.sqrt(V_var)

            #V = V / V_median

            U = []
            m_start = 0
            m_indices = torch.cumsum(ms, dim=0)
            for m_end in m_indices:
                if projector == 'sinkhorn':
                    U.append(self.sinkhorn(V[m_start:m_end, :d], dummy_row=True))
                elif projector == 'hungarian':
                    U.append(self.hungarian(V[m_start:m_end, :d]))
                else:
                    raise NameError('Unknown projector {}.'.format(projector))
                m_start = m_end
            U = torch.cat(U, dim=0)

            #print_helper('iter={}, diff={}, var={}, vmean={}, vvar={}'.format(i, torch.norm(U-lastU), torch.var(torch.sum(U, dim=0)), V_mean, V_var))

            if torch.norm(U - lastU) < 1e-5:
                print_helper(i)
                break

        return U
//  [python]
def forward(self, W, U0, ms, d, projector='sinkhorn'):
        num_graphs = ms.shape[0]

        U = U0
        for i in range(self.max_iter):
            lastU = U
            WU = torch.mm(W, U) #/ num_graphs
            V = torch.linalg.multi_dot([WU, U.t(), WU]) #/ num_graphs ** 2

            #V_median = torch.median(torch.flatten(V, start_dim=-2), dim=-1).values
            #V_var, V_mean = torch.var_mean(torch.flatten(V, start_dim=-2), dim=-1)
            #V = V - V_mean
            #V = V / torch.sqrt(V_var)

            #V = V / V_median

            U = []
            m_start = 0
            m_indices = torch.cumsum(ms, dim=0)
            for m_end in m_indices:
                if projector == 'sinkhorn':
                    U.append(self.sinkhorn(V[m_start:m_end, :d], dummy_row=True))
                elif projector == 'hungarian':
                    U.append(self.hungarian(V[m_start:m_end, :d]))
                else:
                    raise NameError('Unknown projector {}.'.format(projector))
                m_start = m_end
            U = torch.cat(U, dim=0)

            #print_helper('iter={}, diff={}, var={}, vmean={}, vvar={}'.format(i, torch.norm(U-lastU), torch.var(torch.sum(U, dim=0)), V_mean, V_var))

            if torch.norm(U - lastU) < 1e-5:
                print_helper(i)
                break

        return U
//  [python]
def get_weight(self, ):
        u = F.normalize(self.u, dim=1)
        w = self.I - 2 * u.unsqueeze(-1) @ u.unsqueeze(1)
        w = torch.chain_matmul(*[x.squeeze(0) for x in w.chunk(self.num_householders, dim=0)])
        return w.reshape(self.n_dim, 1, self.h, self.w)
//  [python]
def get_weight(self, ):
        u = F.normalize(self.u, dim=1)
        w = self.I - 2 * u.unsqueeze(-1) @ u.unsqueeze(1)
        w = torch.linalg.multi_dot([x.squeeze(0) for x in w.chunk(self.num_householders, dim=0)])
        return w.reshape(self.n_dim, 1, self.h, self.w)
//  [python]
def forward(self):
        out = torch.tensor([]).cuda()
        item_adj = self.item_adj
        x = self.feat
        pre_out = None

        entropy_loss = 0.0
        indepence_loss = 0.0 

        for index in range(len(self.layers)):
            conv_layer = self.conv_layer_list[index]
            x = conv_layer(x, item_adj)

            if self.has_act:
                x = F.leaky_relu(conv_layer(x, item_adj))

            weight = self.weight_list[index]
            temp_out = torch.matmul(x, weight)

            temp_out = torch.softmax(temp_out, dim=1)
            entropy_loss += (-temp_out * torch.log(temp_out+1e-15)).sum(dim=-1).mean()

            item_adj = torch.chain_matmul(temp_out.t(), item_adj, temp_out)

            if pre_out is not None:
                temp_out = torch.matmul(pre_out, temp_out)
            out = torch.cat((out, temp_out), dim=1)
            pre_out = temp_out
            x = weight.t()

            indepence_loss += torch.norm(torch.matmul(weight.t(), weight)-torch.eye(self.layers[index]).cuda(), p=2)/(self.layers[index]*self.layers[index])

        return out, entropy_loss, indepence_loss
//  [python]
def forward(self):
        out = torch.tensor([]).cuda()
        item_adj = self.item_adj
        x = self.feat
        pre_out = None

        entropy_loss = 0.0
        indepence_loss = 0.0 

        for index in range(len(self.layers)):
            conv_layer = self.conv_layer_list[index]
            x = conv_layer(x, item_adj)

            if self.has_act:
                x = F.leaky_relu(conv_layer(x, item_adj))

            weight = self.weight_list[index]
            temp_out = torch.matmul(x, weight)

            temp_out = torch.softmax(temp_out, dim=1)
            entropy_loss += (-temp_out * torch.log(temp_out+1e-15)).sum(dim=-1).mean()

            item_adj = torch.linalg.multi_dot([temp_out.t(), item_adj, temp_out])

            if pre_out is not None:
                temp_out = torch.matmul(pre_out, temp_out)
            out = torch.cat((out, temp_out), dim=1)
            pre_out = temp_out
            x = weight.t()

            indepence_loss += torch.norm(torch.matmul(weight.t(), weight)-torch.eye(self.layers[index]).cuda(), p=2)/(self.layers[index]*self.layers[index])

        return out, entropy_loss, indepence_loss
//  [python]
def test_slogdet_sign(self):
        a = torch.randn(3, 3, requires_grad=True)
        s, logdet = a.slogdet()

        # test that sign should not require grad
        self.assertFalse(s.requires_grad)

        # test that backward through computation involving sign works
        def sign_mul_logdet(mat):
            s, logdet = mat.slogdet()
            return s * logdet

        u, s, v = a.detach().svd()
        s.abs_().clamp_(0.0001)
        for sign in (-1, 1):
            s[-1] = sign
            mat = torch.chain_matmul(u, s.diag(), v.t()).requires_grad_()
            gradcheck(sign_mul_logdet, mat)
            gradgradcheck(sign_mul_logdet, mat)
//  [python]
def test_slogdet_sign(self):
        a = torch.randn(3, 3, requires_grad=True)
        s, logdet = a.slogdet()

        # test that sign should not require grad
        self.assertFalse(s.requires_grad)

        # test that backward through computation involving sign works
        def sign_mul_logdet(mat):
            s, logdet = mat.slogdet()
            return s * logdet

        u, s, v = a.detach().svd()
        s.abs_().clamp_(0.0001)
        for sign in (-1, 1):
            s[-1] = sign
            mat = torch.linalg.multi_dot([u, s.diag(), v.t()]).requires_grad_()
            gradcheck(sign_mul_logdet, mat)
            gradgradcheck(sign_mul_logdet, mat)
//  [python]
def forward(self, inputs, adjacency_matrix):
        '''
        inputs: shape = [num_entity, embedding_dim]
        '''
        outputs = torch.chain_matmul(adjacency_matrix, inputs, self.weights)
        # support = torch.mm(inputs, self.weights)
        # outputs = torch.spmm(inputs, support)
        if self.bias is not None:
            outputs += self.bias
        return outputs
//  [python]
def forward(self, inputs, adjacency_matrix):
        '''
        inputs: shape = [num_entity, embedding_dim]
        '''
        outputs = torch.linalg.multi_dot([adjacency_matrix, inputs, self.weights])
        # support = torch.mm(inputs, self.weights)
        # outputs = torch.spmm(inputs, support)
        if self.bias is not None:
            outputs += self.bias
        return outputs
//  [python]
def compute_G(self, H):
        # the number of hyperedge
        n_edge = H.size(1)
        # the weight of the hyperedge
        we = torch.ones(n_edge).cuda()
        # the degree of the node
        Dv = (H * we).sum(dim=1)
        # the degree of the hyperedge
        De = H.sum(dim=0)

        We = torch.diag(we)
        inv_Dv_half = torch.diag(torch.pow(Dv, -0.5))
        inv_De = torch.diag(torch.pow(De, -1))
        H_T = torch.t(H)

        # propagation matrix
        G = torch.chain_matmul(inv_Dv_half, H, We, inv_De, H_T, inv_Dv_half)

        return G
//  [python]
def compute_G(self, H):
        # the number of hyperedge
        n_edge = H.size(1)
        # the weight of the hyperedge
        we = torch.ones(n_edge).cuda()
        # the degree of the node
        Dv = (H * we).sum(dim=1)
        # the degree of the hyperedge
        De = H.sum(dim=0)

        We = torch.diag(we)
        inv_Dv_half = torch.diag(torch.pow(Dv, -0.5))
        inv_De = torch.diag(torch.pow(De, -1))
        H_T = torch.t(H)

        # propagation matrix
        G = torch.linalg.multi_dot([inv_Dv_half, H, We, inv_De, H_T, inv_Dv_half])

        return G
//  [python]
def _compute_scoring_matrices(self, plda_dim):
        if self.plda_dim != plda_dim:
            self.plda_dim = plda_dim
            #iSt = torch.pinverse(self.St, Settings().backend.plda_rcond)
            #iS = torch.pinverse(self.St - torch.chain_matmul(self.Sb, iSt, self.Sb), Settings().backend.plda_rcond)
            iSt = torch.inverse(self.St)
            iS = torch.inverse(self.St - torch.chain_matmul(self.Sb, iSt, self.Sb))
            Q = iSt - iS
            P = torch.chain_matmul(iSt, self.Sb, iS)
            U, s = torch.svd(P)[:2]
            self.l = s[:plda_dim]
            self.uk = U[:, :plda_dim]
            self.qhat = torch.chain_matmul(self.uk.t(), Q, self.uk)
//  [python]
def _compute_scoring_matrices(self, plda_dim):
        if self.plda_dim != plda_dim:
            self.plda_dim = plda_dim
            iSt = torch.inverse(self.St)
            iS = torch.inverse(self.St - torch.linalg.multi_dot([self.Sb, iSt, self.Sb]))
            Q = iSt - iS
            P = torch.linalg.multi_dot([iSt, self.Sb, iS])
            U, s = torch.svd(P)[:2]
            self.l = s[:plda_dim]
            self.uk = U[:, :plda_dim]
            self.qhat = torch.linalg.multi_dot([self.uk.t(), Q, self.uk])
//  [python]
def _get_weight(self):
        V = self.weights
        V_t = self.weights.transpose(1, 2)
        U = self.I - 2 * torch.bmm(V, V_t) / torch.bmm(V_t, V)
        W = torch.chain_matmul(*U)
        return W
//  [python]
def _get_weight(self):
        V = self.weights
        V_t = self.weights.transpose(1, 2)
        U = self.I - 2 * torch.bmm(V, V_t) / torch.bmm(V_t, V)
        W = torch.linalg.multi_dot(U)
        return W
//  [python]
def test_chain_matmul(self):
        a = torch.randn(3, 4)
        b = torch.randn(4, 5)
        c = torch.randn(5, 6)
        d = torch.randn(6, 7)
        cpu_output = torch.chain_matmul(a, b, c, d)
        npu_output = torch.chain_matmul(a.half().npu(), b.half().npu(), c.half().npu(), d.half().npu())
        self.assertEqual(cpu_output, npu_output, prec=1e-2)
//  [python]
def test_chain_matmul(self):
        a = torch.randn(3, 4)
        b = torch.randn(4, 5)
        c = torch.randn(5, 6)
        d = torch.randn(6, 7)
        cpu_output = torch.linalg.multi_dot([a, b, c, d])
        npu_output = torch.linalg.multi_dot([a.half().npu(), b.half().npu(), c.half().npu(), d.half().npu()])
        self.assertEqual(cpu_output, npu_output, prec=1e-2)
//  [python]
def sample_energy(ceta, mean, cov, zi,n_gmm,bs):
        e = torch.tensor(0.0)
        cov_eps = torch.eye(mean.shape[1]) * (1e-12)
#         cov_eps = cov_eps.to(device)
        for k in range(n_gmm):
            miu_k = mean[k].unsqueeze(1)
            d_k = zi - miu_k

            inv_cov = torch.inverse(cov[k] + cov_eps)
            e_k = torch.exp(-0.5 * torch.chain_matmul(torch.t(d_k), inv_cov, d_k))
            e_k = e_k / torch.sqrt(torch.abs(torch.det(2 * math.pi * cov[k])))
            e_k = e_k * ceta[k]
            e += e_k.squeeze()
        return -torch.log(e)
//  [python]
def sample_energy(ceta, mean, cov, zi,n_gmm,bs):
        e = torch.tensor(0.0)
        cov_eps = torch.eye(mean.shape[1]) * (1e-12)
#         cov_eps = cov_eps.to(device)
        for k in range(n_gmm):
            miu_k = mean[k].unsqueeze(1)
            d_k = zi - miu_k

            inv_cov = torch.inverse(cov[k] + cov_eps)
            e_k = torch.exp(-0.5 * torch.linalg.multi_dot([torch.t(d_k), inv_cov, d_k]))
            e_k = e_k / torch.sqrt(torch.abs(torch.det(2 * math.pi * cov[k])))
            e_k = e_k * ceta[k]
            e += e_k.squeeze()
        return -torch.log(e)
//  [python]
def train_em(cls, data, speaker_labels, plda_dim, iterations, device):
        print('Initializing simplified PLDA...')
        data = data.to(device)     
        n_total_sessions, data_dim = data.size()
        F = torch.randn(data_dim, plda_dim, device=device)
        F = _orthogonalize_columns(F)
        S = 1000 * torch.randn(data_dim, data_dim, device=device)
        data_covariance = torch.matmul(data.t(), data)
        data_list, count_list = _arrange_data_by_counts(data, speaker_labels)
        eye_matrix = torch.eye(plda_dim, device=device)

        for iteration in range(1, iterations+1):
            print('Iteration {}...'.format(iteration), end='')
            iter_start_time = time.time()
            
            FS = torch.solve(F, S.t())[0].t()
            FSF = torch.matmul(FS, F) 
                
            dataEh = torch.zeros(data_dim, plda_dim, device=device)
            Ehh = torch.zeros(plda_dim, plda_dim, device=device)

            for count_data, count in zip(data_list, count_list):
                Sigma = torch.inverse(eye_matrix + count * FSF)
                my = torch.chain_matmul(Sigma, FS.repeat(1, count), count_data.view(-1, data_dim * count).t())            
                dataEh += torch.matmul(count_data.t(), my.repeat(count, 1).t().reshape(count_data.size()[0], -1))
                Ehh += count * (my.size()[1] * Sigma + torch.matmul(my, my.t()))              
            
            F = torch.solve(dataEh.t(), Ehh.t())[0].t()
            S = (data_covariance - torch.chain_matmul(F, Ehh, F.t())) / n_total_sessions

            Sb = torch.matmul(F, F.t())
            St = Sb + S

            print(' [elapsed time = {:0.1f} s]'.format(time.time() - iter_start_time))
            yield Plda(St, Sb)
//  [python]
def train_em(cls, data, speaker_labels, plda_dim, iterations, device):
        print('Initializing simplified PLDA...')
        data = data.to(device)     
        n_total_sessions, data_dim = data.size()
        F = torch.randn(data_dim, plda_dim, device=device)
        F = _orthogonalize_columns(F)
        S = 1000 * torch.randn(data_dim, data_dim, device=device)
        data_covariance = torch.matmul(data.t(), data)
        data_list, count_list = _arrange_data_by_counts(data, speaker_labels)
        eye_matrix = torch.eye(plda_dim, device=device)

        for iteration in range(1, iterations+1):
            print('Iteration {}...'.format(iteration), end='')
            iter_start_time = time.time()
            
            FS = torch.solve(F, S.t())[0].t()
            FSF = torch.matmul(FS, F) 
                
            dataEh = torch.zeros(data_dim, plda_dim, device=device)
            Ehh = torch.zeros(plda_dim, plda_dim, device=device)

            for count_data, count in zip(data_list, count_list):
                Sigma = torch.inverse(eye_matrix + count * FSF)
                my = torch.linalg.multi_dot([Sigma, FS.repeat(1, count), count_data.view(-1, data_dim * count).t()])            
                dataEh += torch.matmul(count_data.t(), my.repeat(count, 1).t().reshape(count_data.size()[0], -1))
                Ehh += count * (my.size()[1] * Sigma + torch.matmul(my, my.t()))              
            
            F = torch.solve(dataEh.t(), Ehh.t())[0].t()
            S = (data_covariance - torch.linalg.multi_dot([F, Ehh, F.t()])) / n_total_sessions

            Sb = torch.matmul(F, F.t())
            St = Sb + S

            print(' [elapsed time = {:0.1f} s]'.format(time.time() - iter_start_time))
            yield Plda(St, Sb)
//  [python]
def _precond_grad_dense_scale(Ql, qr, Grad):
    # type: (Tensor, Tensor, Tensor) -> Tensor
    """
    return preconditioned gradient using (dense, scaling) Kronecker product preconditioner
    Suppose Grad has shape (M, N)
    Ql: shape (M, M), (left side) Cholesky factor of preconditioner
    qr: shape (1, N), defines a diagonal matrix for output feature scaling
    Grad: (matrix) gradient
    """
    return torch.chain_matmul(Ql.t(), Ql, Grad * (qr * qr))
//  [python]
def _precond_grad_dense_scale(Ql, qr, Grad):
    # type: (Tensor, Tensor, Tensor) -> Tensor
    """
    return preconditioned gradient using (dense, scaling) Kronecker product preconditioner
    Suppose Grad has shape (M, N)
    Ql: shape (M, M), (left side) Cholesky factor of preconditioner
    qr: shape (1, N), defines a diagonal matrix for output feature scaling
    Grad: (matrix) gradient
    """
    return torch.linalg.multi_dot([Ql.t(), Ql, Grad * (qr * qr)])
//  [python]
def _precond_grad_norm_dense(ql, Qr, Grad):
    # type: (Tensor, Tensor, Tensor) -> Tensor
    """
    return preconditioned gradient using (normalization, dense) Kronecker product preconditioner
    Suppose Grad has shape (M, N)
    ql[0] is the diagonal part of Ql
    ql[1, 0:-1] is the last column of Ql, excluding the last entry
    Qr: shape (N, N), Cholesky factor of right preconditioner
    Grad: (matrix) gradient
    """
    preG = ql[0:1].t() * Grad + ql[1:].t().mm(Grad[-1:])  # Ql*Grad
    preG = torch.chain_matmul(preG, Qr.t(), Qr)
    add_last_row = ql[1:].mm(preG)  # use it to modify the last row
    preG *= ql[0:1].t()
    preG[-1:] += add_last_row

    return preG
//  [python]
def _precond_grad_norm_dense(ql, Qr, Grad):
    # type: (Tensor, Tensor, Tensor) -> Tensor
    """
    return preconditioned gradient using (normalization, dense) Kronecker product preconditioner
    Suppose Grad has shape (M, N)
    ql[0] is the diagonal part of Ql
    ql[1, 0:-1] is the last column of Ql, excluding the last entry
    Qr: shape (N, N), Cholesky factor of right preconditioner
    Grad: (matrix) gradient
    """
    preG = ql[0:1].t() * Grad + ql[1:].t().mm(Grad[-1:])  # Ql*Grad
    preG = torch.linalg.multi_dot([preG, Qr.t(), Qr])
    add_last_row = ql[1:].mm(preG)  # use it to modify the last row
    preG *= ql[0:1].t()
    preG[-1:] += add_last_row

    return preG
//  [python]
def warp(
		pos1,
		depth1, intrinsics1, pose1, bbox1,
		depth2, intrinsics2, pose2, bbox2
):
	device = pos1.device

	Z1, pos1, ids = interpolate_depth(pos1, depth1)

	# COLMAP convention
	u1 = pos1[1, :] + bbox1[1] + .5
	v1 = pos1[0, :] + bbox1[0] + .5

	X1 = (u1 - intrinsics1[0, 2]) * (Z1 / intrinsics1[0, 0])
	Y1 = (v1 - intrinsics1[1, 2]) * (Z1 / intrinsics1[1, 1])

	XYZ1_hom = torch.cat([
		X1.view(1, -1),
		Y1.view(1, -1),
		Z1.view(1, -1),
		torch.ones(1, Z1.size(0), device=device)
	], dim=0)
	XYZ2_hom = torch.chain_matmul(pose2, torch.inverse(pose1), XYZ1_hom)
	XYZ2 = XYZ2_hom[: -1, :] / XYZ2_hom[-1, :].view(1, -1)

	uv2_hom = torch.matmul(intrinsics2, XYZ2)
	uv2 = uv2_hom[: -1, :] / uv2_hom[-1, :].view(1, -1)

	u2 = uv2[0, :] - bbox2[1] - .5
	v2 = uv2[1, :] - bbox2[0] - .5
	uv2 = torch.cat([u2.view(1, -1),  v2.view(1, -1)], dim=0)

	annotated_depth, pos2, new_ids = interpolate_depth(uv_to_pos(uv2), depth2)

	ids = ids[new_ids]
	pos1 = pos1[:, new_ids]
	estimated_depth = XYZ2[2, new_ids]

	inlier_mask = torch.abs(estimated_depth - annotated_depth) < 0.05

	ids = ids[inlier_mask]
	if ids.size(0) == 0:
		raise EmptyTensorError

	pos2 = pos2[:, inlier_mask]
	pos1 = pos1[:, inlier_mask]

	return pos1, pos2, ids
//  [python]
def warp(
		pos1,
		depth1, intrinsics1, pose1, bbox1,
		depth2, intrinsics2, pose2, bbox2
):
	device = pos1.device

	Z1, pos1, ids = interpolate_depth(pos1, depth1)

	# COLMAP convention
	u1 = pos1[1, :] + bbox1[1] + .5
	v1 = pos1[0, :] + bbox1[0] + .5

	X1 = (u1 - intrinsics1[0, 2]) * (Z1 / intrinsics1[0, 0])
	Y1 = (v1 - intrinsics1[1, 2]) * (Z1 / intrinsics1[1, 1])

	XYZ1_hom = torch.cat([
		X1.view(1, -1),
		Y1.view(1, -1),
		Z1.view(1, -1),
		torch.ones(1, Z1.size(0), device=device)
	], dim=0)
	XYZ2_hom = torch.linalg.multi_dot([pose2, torch.inverse(pose1), XYZ1_hom])
	XYZ2 = XYZ2_hom[: -1, :] / XYZ2_hom[-1, :].view(1, -1)

	uv2_hom = torch.matmul(intrinsics2, XYZ2)
	uv2 = uv2_hom[: -1, :] / uv2_hom[-1, :].view(1, -1)

	u2 = uv2[0, :] - bbox2[1] - .5
	v2 = uv2[1, :] - bbox2[0] - .5
	uv2 = torch.cat([u2.view(1, -1),  v2.view(1, -1)], dim=0)

	annotated_depth, pos2, new_ids = interpolate_depth(uv_to_pos(uv2), depth2)

	ids = ids[new_ids]
	pos1 = pos1[:, new_ids]
	estimated_depth = XYZ2[2, new_ids]

	inlier_mask = torch.abs(estimated_depth - annotated_depth) < 0.05

	ids = ids[inlier_mask]
	if ids.size(0) == 0:
		raise EmptyTensorError

	pos2 = pos2[:, inlier_mask]
	pos1 = pos1[:, inlier_mask]

	return pos1, pos2, ids
//  [python]
def warp(
		pos1,
		depth1, intrinsics1, pose1, bbox1,
		depth2, intrinsics2, pose2, bbox2
):
	device = pos1.device

	Z1, pos1, ids = interpolate_depth(pos1, depth1)
	# COLMAP convention
	u1 = pos1[1, :] + bbox1[1] + .5
	v1 = pos1[0, :] + bbox1[0] + .5

	X1 = (u1 - intrinsics1[0, 2]) * (Z1 / intrinsics1[0, 0])
	Y1 = (v1 - intrinsics1[1, 2]) * (Z1 / intrinsics1[1, 1])

	XYZ1_hom = torch.cat([
		X1.view(1, -1),
		Y1.view(1, -1),
		Z1.view(1, -1),
		torch.ones(1, Z1.size(0), device=device)
	], dim=0)
	XYZ2_hom = torch.chain_matmul(pose2, torch.inverse(pose1), XYZ1_hom)
	XYZ2 = XYZ2_hom[: -1, :] / XYZ2_hom[-1, :].view(1, -1)

	uv2_hom = torch.matmul(intrinsics2, XYZ2)
	uv2 = uv2_hom[: -1, :] / uv2_hom[-1, :].view(1, -1)

	u2 = uv2[0, :] - bbox2[1] - .5
	v2 = uv2[1, :] - bbox2[0] - .5
	uv2 = torch.cat([u2.view(1, -1),  v2.view(1, -1)], dim=0)

	annotated_depth, pos2, new_ids = interpolate_depth(uv_to_pos(uv2), depth2)

	ids = ids[new_ids]
	pos1 = pos1[:, new_ids]
	estimated_depth = XYZ2[2, new_ids]

	differnce = torch.abs(estimated_depth - annotated_depth)
	inlier_mask = torch.abs(estimated_depth - annotated_depth) < 0.05

	ids = ids[inlier_mask]
	if ids.size(0) == 0:
		raise EmptyTensorError

	pos2 = pos2[:, inlier_mask]
	pos1 = pos1[:, inlier_mask]

	return pos1, pos2, ids
//  [python]
def warp(
		pos1,
		depth1, intrinsics1, pose1, bbox1,
		depth2, intrinsics2, pose2, bbox2
):
	device = pos1.device

	Z1, pos1, ids = interpolate_depth(pos1, depth1)
	# COLMAP convention
	u1 = pos1[1, :] + bbox1[1] + .5
	v1 = pos1[0, :] + bbox1[0] + .5

	X1 = (u1 - intrinsics1[0, 2]) * (Z1 / intrinsics1[0, 0])
	Y1 = (v1 - intrinsics1[1, 2]) * (Z1 / intrinsics1[1, 1])

	XYZ1_hom = torch.cat([
		X1.view(1, -1),
		Y1.view(1, -1),
		Z1.view(1, -1),
		torch.ones(1, Z1.size(0), device=device)
	], dim=0)
	XYZ2_hom = torch.linalg.multi_dot([pose2, torch.inverse(pose1), XYZ1_hom])
	XYZ2 = XYZ2_hom[: -1, :] / XYZ2_hom[-1, :].view(1, -1)

	uv2_hom = torch.matmul(intrinsics2, XYZ2)
	uv2 = uv2_hom[: -1, :] / uv2_hom[-1, :].view(1, -1)

	u2 = uv2[0, :] - bbox2[1] - .5
	v2 = uv2[1, :] - bbox2[0] - .5
	uv2 = torch.cat([u2.view(1, -1),  v2.view(1, -1)], dim=0)

	annotated_depth, pos2, new_ids = interpolate_depth(uv_to_pos(uv2), depth2)

	ids = ids[new_ids]
	pos1 = pos1[:, new_ids]
	estimated_depth = XYZ2[2, new_ids]

	differnce = torch.abs(estimated_depth - annotated_depth)
	inlier_mask = torch.abs(estimated_depth - annotated_depth) < 0.05

	ids = ids[inlier_mask]
	if ids.size(0) == 0:
		raise EmptyTensorError

	pos2 = pos2[:, inlier_mask]
	pos1 = pos1[:, inlier_mask]

	return pos1, pos2, ids
//  [python]
def _nested_python_udf(t1, t2, dst):
        t3 = t1 * t2
        t4 = t1 + t2
        res = rpc.rpc_sync('worker{}'.format(dst), my_py_add, args=(t3, t4))
        return torch.chain_matmul(t1, t2, t3, t4, res)
//  [python]
def _nested_python_udf(t1, t2, dst):
        t3 = t1 * t2
        t4 = t1 + t2
        res = rpc.rpc_sync('worker{}'.format(dst), my_py_add, args=(t3, t4))
        return torch.linalg.multi_dot([t1, t2, t3, t4, res])
//  [python]
def build_int_matrix(horizon, diagonal=0, device="cpu", dtype=torch.float32, order=1, traj_dt=None):
    integrate_matrix = torch.tril(
        torch.ones((horizon, horizon), device=device, dtype=dtype), diagonal=diagonal
    )
    chain_list = [torch.eye(horizon, device=device, dtype=dtype)]
    if traj_dt is None:
        chain_list.extend([integrate_matrix for i in range(order)])
    else:
        diag_dt = torch.diag(traj_dt)

        for _ in range(order):
            chain_list.append(integrate_matrix)
            chain_list.append(diag_dt)
    if len(chain_list) == 1:
        integrate_matrix = chain_list[0]
    elif version.parse(torch.__version__) < version.parse("1.9.0"):
        integrate_matrix = torch.chain_matmul(*chain_list)
    else:
        integrate_matrix = torch.linalg.multi_dot(chain_list)

    return integrate_matrix
//  [python]
def build_int_matrix(horizon, diagonal=0, device="cpu", dtype=torch.float32, order=1, traj_dt=None):
    integrate_matrix = torch.tril(
        torch.ones((horizon, horizon), device=device, dtype=dtype), diagonal=diagonal
    )
    chain_list = [torch.eye(horizon, device=device, dtype=dtype)]
    if traj_dt is None:
        chain_list.extend([integrate_matrix for i in range(order)])
    else:
        diag_dt = torch.diag(traj_dt)

        for _ in range(order):
            chain_list.append(integrate_matrix)
            chain_list.append(diag_dt)
    if len(chain_list) == 1:
        integrate_matrix = chain_list[0]
    else:
        integrate_matrix = torch.linalg.multi_dot(chain_list)

    return integrate_matrix
//  [python]
def _nested_rpc_call_backward_error(t1, t2, dst):
        t1 = t1 * t2
        t2 = t1 + t2
        res = rpc.rpc_sync('worker{}'.format(dst),
                           DistAutogradTest._python_udf_with_backward_error,
                           args=(t1, t2))
        return torch.chain_matmul(t1, t2, res)
//  [python]
def _nested_rpc_call_backward_error(t1, t2, dst):
        t1 = t1 * t2
        t2 = t1 + t2
        res = rpc.rpc_sync('worker{}'.format(dst),
                           DistAutogradTest._python_udf_with_backward_error,
                           args=(t1, t2))
        return torch.linalg.multi_dot([t1, t2, res])
//  [python]
def test_chain_matmul(self, device, dtype):
        # chain_matmul accepts a single input tensor while multi_dot does not
        t = make_tensor((2, 2), dtype=dtype, device=device)
        self.assertEqual(t, torch.chain_matmul(t))
        with self.assertRaisesRegex(RuntimeError, r"chain_matmul\(\): Expected one or more matrices"):
            torch.chain_matmul()

        # chain_matmul expects all tensors to be 2D whereas multi_dot allows the first and last tensors to
        # be either 1D or 2D
        with self.assertRaisesRegex(RuntimeError, r"Tensor dimension is 1, expected 2 instead"):
            torch.chain_matmul(make_tensor(1, dtype=dtype, device=device), make_tensor(1, dtype=dtype, device=device))
//  [python]
def test_chain_matmul(self, device, dtype):
        # chain_matmul accepts a single input tensor while multi_dot does not
        t = make_tensor((2, 2), dtype=dtype, device=device)
        self.assertEqual(t, torch.linalg.multi_dot([t]))
        with self.assertRaisesRegex(RuntimeError, r"chain_matmul\(\): Expected one or more matrices"):
            torch.linalg.multi_dot([])

        # chain_matmul expects all tensors to be 2D whereas multi_dot allows the first and last tensors to
        # be either 1D or 2D
        with self.assertRaisesRegex(RuntimeError, r"Tensor dimension is 1, expected 2 instead"):
            torch.linalg.multi_dot([make_tensor(1, dtype=dtype, device=device), make_tensor(1, dtype=dtype, device=device)])
//  [python]
eigh = torch.linalg.eigh if hasattr(torch, 'linalg') else lambda x: torch.symeig(x, eigenvectors=True)
multi_dot = torch.linalg.multi_dot if hasattr(torch, 'linalg') and hasattr(torch.linalg, 'multi_dot') else \
    lambda args: torch.chain_matmul(*args)
//  [python]
eigh = torch.linalg.eigh if hasattr(torch, 'linalg') else lambda x: torch.symeig(x, eigenvectors=True)
multi_dot = torch.linalg.multi_dot if hasattr(torch, 'linalg') and hasattr(torch.linalg, 'multi_dot') else \
    lambda args: torch.linalg.multi_dot(args)
//  [python]
def T(self, x, d, **kwargs):
        # Enable the ability to use particular parameters when T(x,d) is called
        #   as part of the loss function evaluation. For example, the choice of
        #   'L' can change.
        tau    = kwargs.get('tau', self.tau)
        index  = kwargs.get('index', -1)
        assert index >= 0

        mat1 = torch.chain_matmul(self.A.T, self.W1.T, self.W1, self.A)
        mat2 = torch.matmul(self.A.T, self.W2.T)

        hidden = x - self.gamma[index] * F.linear(x, mat1) + self.gamma[index] * F.linear(d, mat2)
        Tx = shrink(hidden, self.theta[index])

        return Tx
//  [python]
def T(self, x, d, **kwargs):
        # Enable the ability to use particular parameters when T(x,d) is called
        #   as part of the loss function evaluation. For example, the choice of
        #   'L' can change.
        tau    = kwargs.get('tau', self.tau)
        index  = kwargs.get('index', -1)
        assert index >= 0

        mat1 = torch.linalg.multi_dot([self.A.T, self.W1.T, self.W1, self.A])
        mat2 = torch.matmul(self.A.T, self.W2.T)

        hidden = x - self.gamma[index] * F.linear(x, mat1) + self.gamma[index] * F.linear(d, mat2)
        Tx = shrink(hidden, self.theta[index])

        return Tx
//  [python]
def _spectral_clustering(self, x, concentration):
        x = x.view(-1, x.shape[-1])
        N = x.shape[0]
        K = self.K if self.K is not None else int(N * self.sample_factor)
        K = max(K, 1)

        A = torch.mm(x, x.t()).mul_(concentration).exp_()
        # remove self-loop
        not_eye = 1 - torch.eye(A.shape[0], dtype=A.dtype, device=A.device)
        A *= not_eye

        D = A.sum(dim=1)
        d = torch.diag(1 / D.sqrt())
        L = torch.diag(D) - A
        norm_L = torch.chain_matmul(d, L, d)
        eig_val, eig_vec = torch.symeig(norm_L, eigenvectors=True)
        eig_vec = utils.normalize_embedding(eig_vec[:, :K])

        init_labels = torch.randint(
            0, K, (N,), dtype=torch.long, device=x.device)
        kmean_labs = (
            segsort.kmeans_with_initial_labels(
                eig_vec, init_labels, iterations=self.iterations))

        return kmean_labs
//  [python]
def _spectral_clustering(self, x, concentration):
        x = x.view(-1, x.shape[-1])
        N = x.shape[0]
        K = self.K if self.K is not None else int(N * self.sample_factor)
        K = max(K, 1)

        A = torch.mm(x, x.t()).mul_(concentration).exp_()
        # remove self-loop
        not_eye = 1 - torch.eye(A.shape[0], dtype=A.dtype, device=A.device)
        A *= not_eye

        D = A.sum(dim=1)
        d = torch.diag(1 / D.sqrt())
        L = torch.diag(D) - A
        norm_L = torch.linalg.multi_dot([d, L, d])
        eig_val, eig_vec = torch.symeig(norm_L, eigenvectors=True)
        eig_vec = utils.normalize_embedding(eig_vec[:, :K])

        init_labels = torch.randint(
            0, K, (N,), dtype=torch.long, device=x.device)
        kmean_labs = (
            segsort.kmeans_with_initial_labels(
                eig_vec, init_labels, iterations=self.iterations))

        return kmean_labs
//  [python]
def compute_lime_signature(self, model, ref_dataset, lime_dataset, cat=True):
        self.logger.info(f"-> compute_lime_signature")
        datasets = []
        model.to(self.device)
        with torch.no_grad():
            phar = tqdm(range(len(lime_dataset)))
            for i in phar:
                lime_data = lime_dataset[i]
                data = ref_dataset[i]
                inputs = torch.from_numpy(data).permute(0, 3, 1, 2).float()
                outputs = ops.batch_forward(model, inputs, argmax=False).detach().cpu().numpy()
                datasets.append([lime_data, outputs])
                phar.set_description(f"-> [{i}/{len(lime_dataset)}] step4.1: compute_lime_signature...")
                out = outputs

        weights = []
        with torch.no_grad():
            phar = tqdm(datasets)
            for data, label in phar:
                data = torch.from_numpy(data).float().to(self.device)
                label = torch.from_numpy(label).float().to(self.device)
                w = torch.chain_matmul(torch.pinverse(torch.matmul(data.T, data)), data.T, label).to("cpu")
                weights.append(w)
                phar.set_description(f"-> [{i}/{len(datasets)}] step4.2: compute weights...")
        if cat:
            weights = torch.cat(weights).to("cpu")
            return out, weights
        else:
            return out, weights.to("cpu")
//  [python]
def compute_lime_signature(self, model, ref_dataset, lime_dataset, cat=True):
        self.logger.info(f"-> compute_lime_signature")
        datasets = []
        model.to(self.device)
        with torch.no_grad():
            phar = tqdm(range(len(lime_dataset)))
            for i in phar:
                lime_data = lime_dataset[i]
                data = ref_dataset[i]
                inputs = torch.from_numpy(data).permute(0, 3, 1, 2).float()
                outputs = ops.batch_forward(model, inputs, argmax=False).detach().cpu().numpy()
                datasets.append([lime_data, outputs])
                phar.set_description(f"-> [{i}/{len(lime_dataset)}] step4.1: compute_lime_signature...")
                out = outputs

        weights = []
        with torch.no_grad():
            phar = tqdm(datasets)
            for data, label in phar:
                data = torch.from_numpy(data).float().to(self.device)
                label = torch.from_numpy(label).float().to(self.device)
                w = torch.linalg.multi_dot([torch.pinverse(torch.matmul(data.T, data)), data.T, label]).to("cpu")
                weights.append(w)
                phar.set_description(f"-> [{i}/{len(datasets)}] step4.2: compute weights...")
        if cat:
            weights = torch.cat(weights).to("cpu")
            return out, weights
        else:
            return out, weights.to("cpu")
//  [python]
def warp(self, depth1, intrinsics1, world2camera1, bbox1,
                   depth2, intrinsics2, world2camera2, bbox2, threshold=0.05):

        device = self.pos.device
        
        Z1 = depth1[self.pos[:,0].long(), self.pos[:,1].long()]
        inds = torch.arange(0, self.pos.size(0), device=device)
        valid_depth_mask = (Z1 != -1)
        Z1 = Z1[valid_depth_mask]
        inds = inds[valid_depth_mask]

        pos1 = self.pos[inds]
        # inds == pos1
        # COLMAP convention
        u1 = pos1[:,1] + bbox1[1] + .5
        v1 = pos1[:,0] + bbox1[0] + .5

        X1 = (u1 - intrinsics1[0, 2]) * (Z1 / intrinsics1[0, 0])
        Y1 = (v1 - intrinsics1[1, 2]) * (Z1 / intrinsics1[1, 1])

        XYZ1_hom = torch.cat([
            X1.view(1, -1),
            Y1.view(1, -1),
            Z1.view(1, -1),
            torch.ones(1, Z1.size(0), device=device)], dim=0)
        XYZ2_hom = torch.chain_matmul(world2camera2, torch.inverse(world2camera1.cpu()).cuda(), XYZ1_hom)
        XYZ2 = XYZ2_hom[: -1, :] / XYZ2_hom[-1, :].view(1, -1)

        uv2_hom = torch.matmul(intrinsics2, XYZ2)
        uv2 = uv2_hom[: -1, :] / uv2_hom[-1, :].view(1, -1)

        u2 = uv2[0, :] - bbox2[1] - .5
        v2 = uv2[1, :] - bbox2[0] - .5
        uv2 = torch.cat([u2.view(1, -1),  v2.view(1, -1)], dim=0)

        pos2 = uv_to_pos(uv2).long()

        bound_mask = (pos2[:,0] >=0) & (pos2[:,0] < depth2.shape[0]) & (pos2[:,1] >=0) & (pos2[:,1] < depth2.shape[1])
        inds = inds[bound_mask]
        pos1 = pos1[bound_mask]
        XYZ2 = XYZ2[:,bound_mask]
        pos2 = pos2[bound_mask]
        new_inds = torch.arange(0, pos2.size(0), device=device)

        annotated_depth = depth2[pos2[:,0].long(), pos2[:,1].long()]

        valid_depth_mask = (annotated_depth != -1)
        annotated_depth = annotated_depth[valid_depth_mask]
        new_inds = new_inds[valid_depth_mask]

        estimated_depth = XYZ2[2, new_inds]
        inlier_mask = torch.abs(estimated_depth - annotated_depth) < threshold

        # pos1 == pos2 == inds
        inds = inds[new_inds][inlier_mask]
        pos1 = pos1[new_inds][inlier_mask]
        pos2 = pos2[new_inds][inlier_mask]

        if inds.size(0) == 0:
            raise EmptyTensorError

        return pos1, pos2