jiahuimbzuai/CoQuIR huggingface dataset:
//  [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.linalg.multi_dot([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
//  [python]
def forward(self, u_features, i_features):
        if self.apply_drop:
            u_features = self.dropout(u_features)
            i_features = self.dropout(i_features)
        if self.accum == 'stack':
            u_features = u_features.reshape(self.num_relations, -1, self.feature_dim)
            i_features = i_features.reshape(self.num_relations, -1, self.feature_dim)
            num_users = u_features.shape[1]
            num_items = i_features.shape[1]
        else:
            num_users = u_features.shape[0]
            num_items = i_features.shape[0]
            
        for relation in range(self.num_relations):
            q_matrix = torch.sum(self.coefs[relation].unsqueeze(1) * self.basis_matrix, 0)
            q_matrix = q_matrix.reshape(self.feature_dim, self.feature_dim)
            if self.accum == 'stack':
                if relation == 0:
                    out = torch.chain_matmul(
                            u_features[relation], q_matrix, 
                            i_features[relation].t()).unsqueeze(-1)
                else:
                    out = torch.cat((out, torch.chain_matmul(
                        u_features[relation], q_matrix, 
                        i_features[relation].t()).unsqueeze(-1)),dim=2)
            else:
                if relation == 0:
                    out = torch.chain_matmul(
                            u_features, q_matrix, i_features.t()).unsqueeze(-1)
                else:
                    out = torch.cat((out, torch.chain_matmul(
                        u_features, q_matrix, i_features.t()).unsqueeze(-1)),dim=2)

        out = out.reshape(num_users * num_items, -1)

        return out
//  [python]
def forward(self, u_features, i_features):
        if self.apply_drop:
            u_features = self.dropout(u_features)
            i_features = self.dropout(i_features)
        if self.accum == 'stack':
            u_features = u_features.reshape(self.num_relations, -1, self.feature_dim)
            i_features = i_features.reshape(self.num_relations, -1, self.feature_dim)
            num_users = u_features.shape[1]
            num_items = i_features.shape[1]
        else:
            num_users = u_features.shape[0]
            num_items = i_features.shape[0]
            
        for relation in range(self.num_relations):
            q_matrix = torch.sum(self.coefs[relation].unsqueeze(1) * self.basis_matrix, 0)
            q_matrix = q_matrix.reshape(self.feature_dim, self.feature_dim)
            if self.accum == 'stack':
                if relation == 0:
                    out = torch.linalg.multi_dot([u_features[relation], q_matrix, i_features[relation].t()]).unsqueeze(-1)
                else:
                    out = torch.cat((out, torch.linalg.multi_dot([u_features[relation], q_matrix, i_features[relation].t()]).unsqueeze(-1)),dim=2)
            else:
                if relation == 0:
                    out = torch.linalg.multi_dot([u_features, q_matrix, i_features.t()]).unsqueeze(-1)
                else:
                    out = torch.cat((out, torch.linalg.multi_dot([u_features, q_matrix, i_features.t()]).unsqueeze(-1)),dim=2)

        out = out.reshape(num_users * num_items, -1)

        return out
//  [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.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 productory2(factors: torch.Tensor, dim=1) -> torch.Tensor:
    """
    Computes the matrix product of the sequence

    This is x1000 slower than 'productory'
    :param factors: * x n x n
    :param dim: acc
    :return: * x n x n: the result of multiplying all matrices on the dim dimension, according
    to the given order. The result will have one less dimension
    """
    b = factors.size(dim - 1)
    m = factors.size(dim)
    acum = []
    for i in range(b):
        mats = [t.squeeze(0) for t in factors.select(dim=dim-1, index=i).chunk(chunks=m)]  # list of m matrices of n x n
        prod = torch.chain_matmul(*mats)
        acum.append(prod)
    return torch.stack(acum, dim=0)
//  [python]
def productory2(factors: torch.Tensor, dim=1) -> torch.Tensor:
    """
    Computes the matrix product of the sequence

    This is x1000 slower than 'productory'
    :param factors: * x n x n
    :param dim: acc
    :return: * x n x n: the result of multiplying all matrices on the dim dimension, according
    to the given order. The result will have one less dimension
    """
    b = factors.size(dim - 1)
    m = factors.size(dim)
    acum = []
    for i in range(b):
        mats = [t.squeeze(0) for t in factors.select(dim=dim-1, index=i).chunk(chunks=m)]  # list of m matrices of n x n
        prod = torch.linalg.multi_dot(mats)
        acum.append(prod)
    return torch.stack(acum, dim=0)
//  [python]
def warp(
        self,
        coord,
        K1,
        pose1,
        bbox1,
        ratio1,
        K2,
        depth2,
        pose2,
        bbox2,
        ratio2,
        max_shape,
    ):
        u1 = (coord[0, :] + bbox1[1] + 0.5) / ratio1[1]
        v1 = (coord[1, :] + bbox1[0] + 0.5) / ratio1[0]

        X1 = (u1 - K1[0, 2]) * (coord[2, :] / K1[0, 0])
        Y1 = (v1 - K1[1, 2]) * (coord[2, :] / K1[1, 1])

        XYZ1_hom = torch.cat(
            [
                X1.view(1, -1),
                Y1.view(1, -1),
                coord[2, :].view(1, -1),
                torch.ones_like(X1.view(1, -1)),
            ],
            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(K2, XYZ2)
        uv2 = uv2_hom[:-1, :] / uv2_hom[-1, :].view(1, -1)

        u2 = uv2[0, :] * ratio2[1] - bbox2[1] - 0.5
        v2 = uv2[1, :] * ratio2[0] - bbox2[0] - 0.5
        uv2 = torch.cat([u2.view(1, -1), v2.view(1, -1)], dim=0)
        # depth check
        ann_z2, _, ids = self.interpolate_depth(uv2, depth2)
        estimate_depth = XYZ2[2, ids]
        uv2 = uv2[:, ids]
        inlier_mask = torch.abs(estimate_depth -
                                ann_z2.to(estimate_depth.device)) < 1.0
        uv2 = uv2[:, inlier_mask]
        if uv2.shape[1] == 0:
            return uv2, uv2

        bx1 = uv2[0, :].min()
        bx2 = uv2[0, :].max()
        by1 = uv2[1, :].min()
        by2 = uv2[1, :].max()
        if max_shape is not None:
            bx1 = bx1.clamp(min=0, max=max_shape[1])
            by1 = by1.clamp(min=0, max=max_shape[0])
            bx2 = bx2.clamp(min=0, max=max_shape[1])
            by2 = by2.clamp(min=0, max=max_shape[0])
        bboxes = torch.stack([bx1, by1, bx2, by2], dim=-1)
        return uv2, bboxes
//  [python]
def warp(
        self,
        coord,
        K1,
        pose1,
        bbox1,
        ratio1,
        K2,
        depth2,
        pose2,
        bbox2,
        ratio2,
        max_shape,
    ):
        u1 = (coord[0, :] + bbox1[1] + 0.5) / ratio1[1]
        v1 = (coord[1, :] + bbox1[0] + 0.5) / ratio1[0]

        X1 = (u1 - K1[0, 2]) * (coord[2, :] / K1[0, 0])
        Y1 = (v1 - K1[1, 2]) * (coord[2, :] / K1[1, 1])

        XYZ1_hom = torch.cat(
            [
                X1.view(1, -1),
                Y1.view(1, -1),
                coord[2, :].view(1, -1),
                torch.ones_like(X1.view(1, -1)),
            ],
            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(K2, XYZ2)
        uv2 = uv2_hom[:-1, :] / uv2_hom[-1, :].view(1, -1)

        u2 = uv2[0, :] * ratio2[1] - bbox2[1] - 0.5
        v2 = uv2[1, :] * ratio2[0] - bbox2[0] - 0.5
        uv2 = torch.cat([u2.view(1, -1), v2.view(1, -1)], dim=0)
        # depth check
        ann_z2, _, ids = self.interpolate_depth(uv2, depth2)
        estimate_depth = XYZ2[2, ids]
        uv2 = uv2[:, ids]
        inlier_mask = torch.abs(estimate_depth -
                                ann_z2.to(estimate_depth.device)) < 1.0
        uv2 = uv2[:, inlier_mask]
        if uv2.shape[1] == 0:
            return uv2, uv2

        bx1 = uv2[0, :].min()
        bx2 = uv2[0, :].max()
        by1 = uv2[1, :].min()
        by2 = uv2[1, :].max()
        if max_shape is not None:
            bx1 = bx1.clamp(min=0, max=max_shape[1])
            by1 = by1.clamp(min=0, max=max_shape[0])
            bx2 = bx2.clamp(min=0, max=max_shape[1])
            by2 = by2.clamp(min=0, max=max_shape[0])
        bboxes = torch.stack([bx1, by1, bx2, by2], dim=-1)
        return uv2, bboxes
//  [python]
def train_lime_model(datasets, device=torch.device('cuda:0' if torch.cuda.is_available() else 'cpu'), cat=True):
    weights = []
    with torch.no_grad():
        for data, label in datasets:
            data, label = torch.from_numpy(data).float().to(device), torch.from_numpy(label).float().to(device)
            w = torch.chain_matmul(torch.pinverse(torch.matmul(data.T, data)), data.T, label)
            weights.append(w)
    if cat:
        return torch.cat(weights)
    else:
        return weights
//  [python]
def train_lime_model(datasets, device=torch.device('cuda:0' if torch.cuda.is_available() else 'cpu'), cat=True):
    weights = []
    with torch.no_grad():
        for data, label in datasets:
            data, label = torch.from_numpy(data).float().to(device), torch.from_numpy(label).float().to(device)
            w = torch.linalg.multi_dot([torch.pinverse(torch.matmul(data.T, data)), data.T, label])
            weights.append(w)
    if cat:
        return torch.cat(weights)
    else:
        return weights
//  [python]
def test_chain_matmul(self, device, dtype):
        def product(matrices):
            for mat in matrices[1:]:
                matrices[0] = matrices[0].mm(mat)
            return matrices[0]

        def run_test(p):
            matrices = []
            for (pi, pi_1) in zip(p[:-1], p[1:]):
                matrices.append(torch.randn(pi, pi_1, dtype=dtype, device=device))
            self.assertEqual(torch.chain_matmul(*matrices), product(matrices))

        run_test([10, 20, 30, 5])
        run_test([15, 5, 10, 20, 25])
//  [python]
def test_chain_matmul(self, device, dtype):
        def product(matrices):
            for mat in matrices[1:]:
                matrices[0] = matrices[0].mm(mat)
            return matrices[0]

        def run_test(p):
            matrices = []
            for (pi, pi_1) in zip(p[:-1], p[1:]):
                matrices.append(torch.randn(pi, pi_1, dtype=dtype, device=device))
            self.assertEqual(torch.linalg.multi_dot(matrices), product(matrices))

        run_test([10, 20, 30, 5])
        run_test([15, 5, 10, 20, 25])
//  [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)
            #print(count_list)
            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())     
                #print(torch.norm(my[:, 0]))          
                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)
            #print(count_list)
            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()])     
                #print(torch.norm(my[:, 0]))          
                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 forward(self, u_features, i_features):
        if self.apply_drop:
            u_features = self.dropout(u_features)
            i_features = self.dropout(i_features)
        if self.accum == 'stack':
            u_features = u_features.reshape(self.num_relations, -1, self.feature_dim)
            i_features = i_features.reshape(self.num_relations, -1, self.feature_dim)
            num_users = u_features.shape[1]
            num_items = i_features.shape[1]
        else:
            num_users = u_features.shape[0]
            num_items = i_features.shape[0]
        for relation in range(self.num_relations):
            q_matrix = torch.sum(self.coefs[relation].unsqueeze(1) * self.basis_matrix, 0)
            q_matrix = q_matrix.reshape(self.feature_dim, self.feature_dim)
            if self.accum == 'stack':
                if relation == 0:
                    out = torch.chain_matmul(u_features[relation], q_matrix, i_features[relation].t()).unsqueeze(-1)
                else:
                    out = torch.cat((out, torch.chain_matmul(u_features[relation], q_matrix, i_features[relation].t()).unsqueeze(-1)), dim=2)
            elif relation == 0:
                out = torch.chain_matmul(u_features, q_matrix, i_features.t()).unsqueeze(-1)
            else:
                out = torch.cat((out, torch.chain_matmul(u_features, q_matrix, i_features.t()).unsqueeze(-1)), dim=2)
        out = out.reshape(num_users * num_items, -1)
        return out
//  [python]
def forward(self, u_features, i_features):
        if self.apply_drop:
            u_features = self.dropout(u_features)
            i_features = self.dropout(i_features)
        if self.accum == 'stack':
            u_features = u_features.reshape(self.num_relations, -1, self.feature_dim)
            i_features = i_features.reshape(self.num_relations, -1, self.feature_dim)
            num_users = u_features.shape[1]
            num_items = i_features.shape[1]
        else:
            num_users = u_features.shape[0]
            num_items = i_features.shape[0]
        for relation in range(self.num_relations):
            q_matrix = torch.sum(self.coefs[relation].unsqueeze(1) * self.basis_matrix, 0)
            q_matrix = q_matrix.reshape(self.feature_dim, self.feature_dim)
            if self.accum == 'stack':
                if relation == 0:
                    out = torch.linalg.multi_dot([u_features[relation], q_matrix, i_features[relation].t()]).unsqueeze(-1)
                else:
                    out = torch.cat((out, torch.linalg.multi_dot([u_features[relation], q_matrix, i_features[relation].t()]).unsqueeze(-1)), dim=2)
            elif relation == 0:
                out = torch.linalg.multi_dot([u_features, q_matrix, i_features.t()]).unsqueeze(-1)
            else:
                out = torch.cat((out, torch.linalg.multi_dot([u_features, q_matrix, i_features.t()]).unsqueeze(-1)), dim=2)
        out = out.reshape(num_users * num_items, -1)
        return out
//  [python]
def forward(
        self,        
        depth1, intrinsics1, pose1, bbox1,
        depth2, intrinsics2, pose2, bbox2):
        
        # depth1 = depth1.view(1,1,self.size,self.size)
        # depth1 = KF.gaussian_blur2d(depth1,kernel_size=(15,15),sigma=(5,5)).view(-1)
        Z1 = depth1.view(-1)
        u1 = self.raw_pos[1, :] + bbox1[1] + .5
        v1 = self.raw_pos[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),
            self.ones
        ], 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
        u2 = u2.long()
        v2 = v2.long()
        valid_pos2 = torch.logical_and(u2>=0,u2<depth1.shape[1])
        valid_pos2 = torch.logical_and(valid_pos2,v2>=0)
        valid_pos2 = torch.logical_and(valid_pos2,v2<depth1.shape[0])
        outlier = torch.logical_not(valid_pos2)
        v2[outlier] = 0
        u2[outlier] = 0
        annotated_depth = depth2[v2,u2]
        estimated_depth = XYZ2[2]
        annotated_depth = annotated_depth.view(1,1,self.size,self.size)
        annotated_depth = KF.gaussian_blur2d(annotated_depth,kernel_size=(15,15),sigma=(5,5))
        estimated_depth = estimated_depth.view(1,1,self.size,self.size)
        estimated_depth = KF.gaussian_blur2d(estimated_depth,kernel_size=(15,15),sigma=(5,5))
        # print(outlier.float().sum())
        # t = torch.abs(annotated_depth - estimated_depth) > 0.05
        outlier = outlier.view(1,1,self.size,self.size)
        outlier = torch.logical_or(torch.abs(annotated_depth - estimated_depth) > 0.05,outlier)
        # outlier =  KF.gaussian_blur2d(outlier.float(),kernel_size=(15,15),sigma=(5,5))>0.9
        outlier = outlier.view(-1)
        # if outlier.float().sum()>self.size*self.size*0.8:
        #     print('error')
        u2[outlier] = 1e9
        v2[outlier] = 1e9
        flow = torch.cat([u2.view(1,self.size,self.size,1),v2.view(1,self.size,self.size,1)],dim=-1)

        return flow
//  [python]
def forward(
        self,        
        depth1, intrinsics1, pose1, bbox1,
        depth2, intrinsics2, pose2, bbox2):
        
        # depth1 = depth1.view(1,1,self.size,self.size)
        # depth1 = KF.gaussian_blur2d(depth1,kernel_size=(15,15),sigma=(5,5)).view(-1)
        Z1 = depth1.view(-1)
        u1 = self.raw_pos[1, :] + bbox1[1] + .5
        v1 = self.raw_pos[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),
            self.ones
        ], 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
        u2 = u2.long()
        v2 = v2.long()
        valid_pos2 = torch.logical_and(u2>=0,u2<depth1.shape[1])
        valid_pos2 = torch.logical_and(valid_pos2,v2>=0)
        valid_pos2 = torch.logical_and(valid_pos2,v2<depth1.shape[0])
        outlier = torch.logical_not(valid_pos2)
        v2[outlier] = 0
        u2[outlier] = 0
        annotated_depth = depth2[v2,u2]
        estimated_depth = XYZ2[2]
        annotated_depth = annotated_depth.view(1,1,self.size,self.size)
        annotated_depth = KF.gaussian_blur2d(annotated_depth,kernel_size=(15,15),sigma=(5,5))
        estimated_depth = estimated_depth.view(1,1,self.size,self.size)
        estimated_depth = KF.gaussian_blur2d(estimated_depth,kernel_size=(15,15),sigma=(5,5))
        # print(outlier.float().sum())
        # t = torch.abs(annotated_depth - estimated_depth) > 0.05
        outlier = outlier.view(1,1,self.size,self.size)
        outlier = torch.logical_or(torch.abs(annotated_depth - estimated_depth) > 0.05,outlier)
        # outlier =  KF.gaussian_blur2d(outlier.float(),kernel_size=(15,15),sigma=(5,5))>0.9
        outlier = outlier.view(-1)
        # if outlier.float().sum()>self.size*self.size*0.8:
        #     print('error')
        u2[outlier] = 1e9
        v2[outlier] = 1e9
        flow = torch.cat([u2.view(1,self.size,self.size,1),v2.view(1,self.size,self.size,1)],dim=-1)

        return flow
//  [python]
def _get_credit(self, user_idx, X, tilde_X):
        # Please refer to Eq.8.
        Minv, tilde_Minv = self.Minv[user_idx], self.tilde_Minv[user_idx]

        tilde_Minv_tilde_X = torch.mm(tilde_Minv, tilde_X.T)
        result_a = torch.chain_matmul(X, Minv, tilde_Minv_tilde_X)
        result_b = 1 + (tilde_X.T * tilde_Minv_tilde_X).sum(dim=0)
        norm_M = result_a.norm(dim=0)

        return norm_M * norm_M / result_b
//  [python]
def _get_credit(self, user_idx, X, tilde_X):
        # Please refer to Eq.8.
        Minv, tilde_Minv = self.Minv[user_idx], self.tilde_Minv[user_idx]

        tilde_Minv_tilde_X = torch.mm(tilde_Minv, tilde_X.T)
        result_a = torch.linalg.multi_dot([X, Minv, tilde_Minv_tilde_X])
        result_b = 1 + (tilde_X.T * tilde_Minv_tilde_X).sum(dim=0)
        norm_M = result_a.norm(dim=0)

        return norm_M * norm_M / result_b
//  [python]
def _update(self, s, y, rho_inv):
        rho = rho_inv.reciprocal()
        if self.inverse:
            if self.n_updates == 0:
                self.H.mul_(rho_inv / y.dot(y))
            torch.addr(
                torch.chain_matmul(
                    torch.addr(self.I, s, y, alpha=-rho),
                    self.H,
                    torch.addr(self.I, y, s, alpha=-rho),
                ),
                s,
                s,
                alpha=rho,
                out=self.H,
            )
        else:
            if self.n_updates == 0:
                self.B.mul_(rho * y.dot(y))
            Bs = torch.mv(self.B, s)
            torch.addr(
                torch.addr(self.B, y, y, alpha=rho),
                Bs,
                Bs,
                alpha=s.dot(Bs).reciprocal().neg(),
                out=self.B,
            )
//  [python]
def _update(self, s, y, rho_inv):
        rho = rho_inv.reciprocal()
        if self.inverse:
            if self.n_updates == 0:
                self.H.mul_(rho_inv / y.dot(y))
            torch.addr(
                torch.linalg.multi_dot(
                    [torch.addr(self.I, s, y, alpha=-rho),
                     self.H,
                     torch.addr(self.I, y, s, alpha=-rho),]
                ),
                s,
                s,
                alpha=rho,
                out=self.H,
            )
        else:
            if self.n_updates == 0:
                self.B.mul_(rho * y.dot(y))
            Bs = torch.mv(self.B, s)
            torch.addr(
                torch.addr(self.B, y, y, alpha=rho),
                Bs,
                Bs,
                alpha=s.dot(Bs).reciprocal().neg(),
                out=self.B,
            )
//  [python]
def forward(self, inputs,edge_inputs, adjacency_matrix):
		'''
		inputs: shape = [num_entity, embedding_dim]
		'''
		outputs_edge = torch.einsum('ijk,kp->ijp',edge_inputs,self.weights_edge)
		outputs_edge = torch.mean(outputs_edge,dim = 1)
		outputs_node = torch.chain_matmul(adjacency_matrix, inputs, self.weights_node)
		outputs = outputs_edge + outputs_node

		if self.bias is not None:
			outputs += self.bias
		node_weight = torch.sum(adjacency_matrix,1)
		node_weight_zero = torch.eq(node_weight,0).float()
		node_weight = (node_weight + node_weight_zero).unsqueeze(1).expand_as(outputs)
		return outputs / node_weight
//  [python]
def forward(self, inputs,edge_inputs, adjacency_matrix):
		'''
		inputs: shape = [num_entity, embedding_dim]
		'''
		outputs_edge = torch.einsum('ijk,kp->ijp',edge_inputs,self.weights_edge)
		outputs_edge = torch.mean(outputs_edge,dim = 1)
		outputs_node = torch.linalg.multi_dot([adjacency_matrix, inputs, self.weights_node])
		outputs = outputs_edge + outputs_node

		if self.bias is not None:
			outputs += self.bias
		node_weight = torch.sum(adjacency_matrix,1)
		node_weight_zero = torch.eq(node_weight,0).float()
		node_weight = (node_weight + node_weight_zero).unsqueeze(1).expand_as(outputs)
		return outputs / node_weight
//  [python]
def forward(self, inputs,edge_inputs, adjacency_matrix):
		'''
		inputs: shape = [num_entity, embedding_dim]
		'''
		outputs_edge = torch.einsum('ijk,kp->ijp',edge_inputs,self.weights_edge)
		#outputs_edge = self.linear_edge(edge_inputs)
		#print ('weight_edge',self.weights_edge)
		#outputs_edge = torch.bmm(edge_inputs,self.weights_edge.unsqueeze(0).expand(self.input_dim))
		outputs_edge = torch.mean(outputs_edge,dim = 1)
		outputs_node = torch.chain_matmul(adjacency_matrix, inputs, self.weights_node)
		outputs = outputs_edge + outputs_node
		# support = torch.mm(inputs, self.weights)
		# outputs = torch.spmm(inputs, support)
		if self.bias is not None:
			outputs += self.bias
		node_weight = torch.sum(adjacency_matrix,1)
		node_weight_zero = torch.eq(node_weight,0).float()
		node_weight = (node_weight + node_weight_zero).unsqueeze(1).expand_as(outputs)
		return outputs / node_weight
//  [python]
def forward(self, inputs,edge_inputs, adjacency_matrix):
		'''
		inputs: shape = [num_entity, embedding_dim]
		'''
		outputs_edge = torch.einsum('ijk,kp->ijp',edge_inputs,self.weights_edge)
		#outputs_edge = self.linear_edge(edge_inputs)
		#print ('weight_edge',self.weights_edge)
		#outputs_edge = torch.bmm(edge_inputs,self.weights_edge.unsqueeze(0).expand(self.input_dim))
		outputs_edge = torch.mean(outputs_edge,dim = 1)
		outputs_node = torch.linalg.multi_dot([adjacency_matrix, inputs, self.weights_node])
		outputs = outputs_edge + outputs_node
		# support = torch.mm(inputs, self.weights)
		# outputs = torch.spmm(inputs, support)
		if self.bias is not None:
			outputs += self.bias
		node_weight = torch.sum(adjacency_matrix,1)
		node_weight_zero = torch.eq(node_weight,0).float()
		node_weight = (node_weight + node_weight_zero).unsqueeze(1).expand_as(outputs)
		return outputs / node_weight
//  [python]
def point_vector_angle(c, p0, p1):
    #all inputs are torch.Tensor
    #"c" is center point - shape=(number of dimenisons)
    #"p0" is the point in the direction of zero degrees - shape=(number of dimenisons)
    #"p1" are the points to find an angle to - shape=(number of dimenisons, number of points)
    #outputs the angle between the vectors in degrees
    if len(p1.shape)<2: p1 = p1[:,None]
    v0 = (p0-c)
    v1 = p1-c[:,None]
    tmp = torch.sum(v0*v1, dim=0)/(torch.norm(v0)*torch.norm(v1,dim=0))
    return torch.acos(tmp)/np.pi*180
//  [python]
def point_vector_angle(c, p0, p1):
    #all inputs are torch.Tensor
    #"c" is center point - shape=(number of dimenisons)
    #"p0" is the point in the direction of zero degrees - shape=(number of dimenisons)
    #"p1" are the points to find an angle to - shape=(number of dimenisons, number of points)
    #outputs the angle between the vectors in degrees
    if len(p1.shape)<2: p1 = p1[:,None]
    v0 = (p0-c)
    v1 = p1-c[:,None]
    tmp = torch.linalg.multi_dot([v0,v1])/(torch.norm(v0)*torch.norm(v1,dim=0))
    return torch.acos(tmp)/np.pi*180
//  [python]
def _linear_raw_residual_from_jac(linear_jac, mesh, sol, forc_fun, auto_jac):
        res = t.matmul(linear_jac, sol)
        res += forc_fun(mesh.mesh_pts)[:, 0]
        return res
//  [python]
def _linear_raw_residual_from_jac(linear_jac, mesh, sol, forc_fun, auto_jac):
        res = multi_dot((linear_jac, sol))
        res += forc_fun(mesh.mesh_pts)[:, 0]
        return res
//  [python]
def _linear_solve(self, jac, res):
        if self._islinear:
            if self._linear_jac_factors is None:
                import numpy as np
                self._linear_jac_factors = np.linalg.qr(
                    jac.numpy(), mode="complete")
                self._linear_jac_factors = [
                    torch.as_tensor(f, dtype=torch.double)
                    for f in self._linear_jac_factors]
            tmp = torch.matmul(self._linear_jac_factors[0].T, res)[:, None]
            try:
                return torch.linalg.solve_triangular(
                    self._linear_jac_factors[1], tmp, upper=True)[:, 0]
            except AttributeError:
                return torch.triangular_solve(
                    tmp, self._linear_jac_factors[1], upper=True)[0][:, 0]

        return super()._linear_solve(jac, res)
//  [python]
def _linear_solve(self, jac, res):
        if self._islinear:
            if self._linear_jac_factors is None:
                # this line does not work with pickle
                # self._linear_jac_factors = torch.linalg.qr(
                #     jac, mode="complete")
                # so use these three lines
                import numpy as np
                self._linear_jac_factors = np.linalg.qr(
                    jac.numpy(), mode="complete")
                self._linear_jac_factors = [
                    torch.as_tensor(f, dtype=torch.double)
                    for f in self._linear_jac_factors]
            # print((torch.linalg.multi_dot(self._linear_jac_factors)))
            # print(jac)
            # print((torch.linalg.multi_dot(self._linear_jac_factors))-jac)
            # assert torch.allclose(
            #     torch.linalg.multi_dot(self._linear_jac_factors), jac)

            tmp = torch.linalg.multi_dot(
                (self._linear_jac_factors[0].T, res))[:, None]
            try:
                return torch.linalg.solve_triangular(
                    self._linear_jac_factors[1], tmp, upper=True)[:, 0]
            except AttributeError:
                return torch.triangular_solve(
                    tmp, self._linear_jac_factors[1], upper=True)[0][:, 0]

        return super()._linear_solve(jac, res)
//  [python]
def _linear_raw_residual_jac(mesh, diff_fun, vel_fun):
        diff_vals = diff_fun(mesh.mesh_pts)
        # assert torch.all(diff_vals > 0)
        vel_vals = vel_fun(mesh.mesh_pts)
        linear_jac = 0
        for dd in range(mesh.nphys_vars):
            linear_jac += (
                torch.matmul(mesh._dmats[dd], torch.matmul(diff_vals*mesh._dmats[dd], mesh._dmats[dd].T)) -
                vel_vals[:, dd:dd+1]*mesh._dmats[dd])
        return linear_jac
//  [python]
def _linear_raw_residual_jac(mesh, diff_fun, vel_fun):
        diff_vals = diff_fun(mesh.mesh_pts)
        # assert torch.all(diff_vals > 0)
        vel_vals = vel_fun(mesh.mesh_pts)
        linear_jac = 0
        for dd in range(mesh.nphys_vars):
            linear_jac += (
                multi_dot(
                    (mesh._dmats[dd], diff_vals*mesh._dmats[dd])) -
                vel_vals[:, dd:dd+1]*mesh._dmats[dd])
        return linear_jac
//  [python]
def sym_inverse(mat, epsilon: float = 1e-6, return_sqrt=False, mode='regularize', return_both=False):
    """ Utility function that returns the inverse of a matrix, with the
    option to return the square root of the inverse matrix.

    Parameters
    ----------
    mat: numpy array with shape [m,m]
        Matrix to be inverted.
    epsilon : float
        Cutoff for eigenvalues.
    return_sqrt: bool, optional, default = False
        if True, the square root of the inverse matrix is returned instead
    mode: str, default='trunc'
        Whether to truncate eigenvalues if they are too small or to regularize them by taking the absolute value
        and adding a small positive constant. :code:`trunc` leads to truncation, :code:`regularize` leads to epsilon
        being added to the eigenvalues after taking the absolute value
    return_both: bool, default=False
        Whether to return the sqrt and its inverse or simply the inverse
    Returns
    -------
    x_inv: numpy array with shape [m,m]
        inverse of the original matrix
    """
    eigval, eigvec = symeig_reg(mat, epsilon, mode)

    # Build the diagonal matrix with the filtered eigenvalues or square
    # root of the filtered eigenvalues according to the parameter
    if return_sqrt:
        diag_inv = torch.diag(torch.sqrt(1. / eigval))
        if return_both:
            diag = torch.diag(torch.sqrt(eigval))
    else:
        diag_inv = torch.diag(1. / eigval)
        if return_both:
            diag = torch.diag(eigval)
    if not return_both:
        return torch.chain_matmul(eigvec.t(), diag_inv, eigvec)
    else:
        return torch.chain_matmul(eigvec.t(), diag_inv, eigvec), torch.chain_matmul(eigvec.t(), diag, eigvec)
//  [python]
def sym_inverse(mat, epsilon: float = 1e-6, return_sqrt=False, mode='regularize', return_both=False):
    """ Utility function that returns the inverse of a matrix, with the
    option to return the square root of the inverse matrix.

    Parameters
    ----------
    mat: numpy array with shape [m,m]
        Matrix to be inverted.
    epsilon : float
        Cutoff for eigenvalues.
    return_sqrt: bool, optional, default = False
        if True, the square root of the inverse matrix is returned instead
    mode: str, default='trunc'
        Whether to truncate eigenvalues if they are too small or to regularize them by taking the absolute value
        and adding a small positive constant. :code:`trunc` leads to truncation, :code:`regularize` leads to epsilon
        being added to the eigenvalues after taking the absolute value
    return_both: bool, default=False
        Whether to return the sqrt and its inverse or simply the inverse
    Returns
    -------
    x_inv: numpy array with shape [m,m]
        inverse of the original matrix
    """
    eigval, eigvec = symeig_reg(mat, epsilon, mode)

    # Build the diagonal matrix with the filtered eigenvalues or square
    # root of the filtered eigenvalues according to the parameter
    if return_sqrt:
        diag_inv = torch.diag(torch.sqrt(1. / eigval))
        if return_both:
            diag = torch.diag(torch.sqrt(eigval))
    else:
        diag_inv = torch.diag(1. / eigval)
        if return_both:
            diag = torch.diag(eigval)
    if not return_both:
        return torch.linalg.multi_dot([eigvec.t(), diag_inv, eigvec])
    else:
        return torch.linalg.multi_dot([eigvec.t(), diag_inv, eigvec]), torch.linalg.multi_dot([eigvec.t(), diag, eigvec])
//  [python]
def multidot(H,y):
    ret = torch.matmul(torch.matmul(torch.t(y), torch.t(H)), H)
    return ret
//  [python]
def multidot(H,y):
    ret = torch.linalg.multi_dot([torch.t(y), torch.t(H), H]) 
    return ret
//  [python]
def partial_deriv(deriv_mats, quantity, dd, idx=None):
    if idx is None:
        return torch.mm(*[deriv_mats[dd], quantity])
    return torch.mm(*[deriv_mats[dd][idx], quantity])
//  [python]
def partial_deriv(deriv_mats, quantity, dd, idx=None):
    if idx is None:
        return torch.linalg.multi_dot((deriv_mats[dd], quantity))
    return torch.linalg.multi_dot((deriv_mats[dd][idx], quantity))
//  [python]
def high_order_partial_deriv(order, deriv_mats, quantity, dd, idx=None):
    Dmat = torch.mm(*[deriv_mats[dd] for ii in range(order)])
    if idx is None:
        return torch.mm(Dmat, quantity)
    return torch.mm(Dmat[idx], quantity)
//  [python]
def high_order_partial_deriv(order, deriv_mats, quantity, dd, idx=None):
    Dmat = torch.linalg.multi_dot([deriv_mats[dd] for ii in range(order)])
    if idx is None:
        return torch.linalg.multi_dot((Dmat, quantity))
    return torch.linalg.multi_dot((Dmat[idx], quantity))
//  [python]
def _solve_v_and_rescale(self, weight_mat, u, target_sigma):
        # Tries to returns a vector `v` s.t. `u = normalize(W @ v)`
        # (the invariant at top of this class) and `u @ W @ v = sigma`.
        # This uses pinverse in case W^T W is not invertible.
        v = torch.mm(torch.mm(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):
        # Tries to returns a vector `v` s.t. `u = normalize(W @ v)`
        # (the invariant at top of this class) and `u @ W @ v = sigma`.
        # This uses pinverse in case W^T W is not invertible.
        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 blas_lapack_ops(self):
        m = torch.randn(3, 3)
        a = torch.randn(10, 3, 4)
        b = torch.randn(10, 4, 3)
        v = torch.randn(3)
        return len(
            torch.addbmm(m, a, b),
            torch.addmm(torch.randn(2, 3), torch.randn(2, 3), torch.randn(3, 3)),
            torch.addmv(torch.randn(2), torch.randn(2, 3), torch.randn(3)),
            torch.addr(torch.zeros(3, 3), v, v),
            torch.baddbmm(m, a, b),
            torch.bmm(a, b),
            torch.chain_matmul(torch.randn(3, 3), torch.randn(3, 3), torch.randn(3, 3)),
            # torch.cholesky(a), # deprecated
            # torch.cholesky_inverse(torch.randn(3, 3)), # had some error
            # torch.cholesky_solve(torch.randn(3, 3), torch.randn(3, 3)),
            torch.dot(v, v),
            # torch.linalg.eig(m), # not build with lapack
            # torch.geqrf(a),
            torch.ger(v, v),
            torch.inner(m, m),
            # torch.inverse(m),
            # torch.det(m),
            # torch.logdet(m),
            # torch.slogdet(m),
            # torch.lstsq(m, m),
            # torch.linalg.lu_factor(m),
            # torch.lu_solve(m, *torch.linalg.lu_factor(m)),
            # torch.lu_unpack(*torch.linalg.lu_factor(m)),
            torch.matmul(m, m),
            torch.matrix_power(m, 2),
            # torch.matrix_rank(m),
            torch.matrix_exp(m),
            torch.mm(m, m),
            torch.mv(m, v),
            # torch.orgqr(a, m),
            # torch.ormqr(a, m, v),
            torch.outer(v, v),
            # torch.pinverse(m),
            # torch.qr(a),
            # torch.solve(m, m),
            # torch.svd(a),
            # torch.svd_lowrank(a),
            # torch.pca_lowrank(a),
            # torch.symeig(a), # deprecated
            # torch.lobpcg(a, b), # not supported
            torch.trapz(m, m),
            torch.trapezoid(m, m),
            torch.cumulative_trapezoid(m, m),
            # torch.triangular_solve(m, m),
            torch.vdot(v, v),
        )
//  [python]
def blas_lapack_ops(self):
        m = torch.randn(3, 3)
        a = torch.randn(10, 3, 4)
        b = torch.randn(10, 4, 3)
        v = torch.randn(3)
        return len(
            torch.addbmm(m, a, b),
            torch.addmm(torch.randn(2, 3), torch.randn(2, 3), torch.randn(3, 3)),
            torch.addmv(torch.randn(2), torch.randn(2, 3), torch.randn(3)),
            torch.addr(torch.zeros(3, 3), v, v),
            torch.baddbmm(m, a, b),
            torch.bmm(a, b),
            torch.linalg.multi_dot([torch.randn(3, 3), torch.randn(3, 3), torch.randn(3, 3)]),
            # torch.cholesky(a), # deprecated
            # torch.cholesky_inverse(torch.randn(3, 3)), # had some error
            # torch.cholesky_solve(torch.randn(3, 3), torch.randn(3, 3)),
            torch.dot(v, v),
            # torch.linalg.eig(m), # not build with lapack
            # torch.geqrf(a),
            torch.ger(v, v),
            torch.inner(m, m),
            # torch.inverse(m),
            # torch.det(m),
            # torch.logdet(m),
            # torch.slogdet(m),
            # torch.lstsq(m, m),
            # torch.linalg.lu_factor(m),
            # torch.lu_solve(m, *torch.linalg.lu_factor(m)),
            # torch.lu_unpack(*torch.linalg.lu_factor(m)),
            torch.matmul(m, m),
            torch.matrix_power(m, 2),
            # torch.matrix_rank(m),
            torch.matrix_exp(m),
            torch.mm(m, m),
            torch.mv(m, v),
            # torch.orgqr(a, m),
            # torch.ormqr(a, m, v),
            torch.outer(v, v),
            # torch.pinverse(m),
            # torch.qr(a),
            # torch.solve(m, m),
            # torch.svd(a),
            # torch.svd_lowrank(a),
            # torch.pca_lowrank(a),
            # torch.symeig(a), # deprecated
            # torch.lobpcg(a, b), # not supported
            torch.trapz(m, m),
            torch.trapezoid(m, m),
            torch.cumulative_trapezoid(m, m),
            # torch.triangular_solve(m, m),
            torch.vdot(v, v),
        )
//  [python]
def _raw_jacobian_nD(self, split_sols, depth_vals):
        dudx_ij = self._derivs(split_sols)
        visc = self._viscosity(dudx_ij)
        visc_jac = self._viscosity_jac(dudx_ij)
        vecs = self._vector_components(dudx_ij)
        vecs_jac = self._vector_components_jac(dudx_ij)
        # dmats = [self.mesh._meshes[0]._dmat(dd)
        #          for dd in range(self.mesh.nphys_vars)]
        dmats = self.mesh._meshes[0]._dmats
        jac = []
        # loop over jacobian rows (blocks)
        for ii in range(self.mesh.nphys_vars):
            jac.append([0 for jj in range(self.mesh.nphys_vars)])
            # loop over jacobian columns (blocks)
            for jj in range(self.mesh.nphys_vars):
                # loop over phys_vars to compute divergence
                for dd in range(self.mesh.nphys_vars):
                    jac[ii][jj] += torch.matmul(dmats[dd], (
                        (2*depth_vals[:, :1]*vecs[ii][:, dd:dd+1]*visc_jac[jj]) +
                        ((2*visc*depth_vals[:, 0])[:, None])*vecs_jac[ii][jj][dd]))
            jac[ii][ii] -= torch.diag(self._beta_vals[:, 0])
        return torch.vstack([torch.hstack(j) for j in jac])
//  [python]
def _raw_jacobian_nD(self, split_sols, depth_vals):
        dudx_ij = self._derivs(split_sols)
        visc = self._viscosity(dudx_ij)
        visc_jac = self._viscosity_jac(dudx_ij)
        vecs = self._vector_components(dudx_ij)
        vecs_jac = self._vector_components_jac(dudx_ij)
        # dmats = [self.mesh._meshes[0]._dmat(dd)
        #          for dd in range(self.mesh.nphys_vars)]
        dmats = self.mesh._meshes[0]._dmats
        jac = []
        # loop over jacobian rows (blocks)
        for ii in range(self.mesh.nphys_vars):
            jac.append([0 for jj in range(self.mesh.nphys_vars)])
            # loop over jacobian columns (blocks)
            for jj in range(self.mesh.nphys_vars):
                # loop over phys_vars to compute divergence
                for dd in range(self.mesh.nphys_vars):
                    jac[ii][jj] += multi_dot((dmats[dd], (
                        (2*depth_vals[:, :1]*vecs[ii][:, dd:dd+1]*visc_jac[jj]) +
                        ((2*visc*depth_vals[:, 0])[:, None])*vecs_jac[ii][jj][dd])))
            jac[ii][ii] -= torch.diag(self._beta_vals[:, 0])
        return torch.vstack([torch.hstack(j) for j in jac])
//  [python]
def _cheby_interpolate(self, canonical_abscissa_1d,
                           canonical_barycentric_weights_1d, values,
                           canonical_eval_samples):
        # if type(values) != np.ndarray:
        #     values = values.detach().numpy()
        # interp_vals = multivariate_barycentric_lagrange_interpolation(
        #     canonical_eval_samples, canonical_abscissa_1d,
        #     canonical_barycentric_weights_1d, values,
        #     np.arange(self.nphys_vars))
        values = torch.as_tensor(values, dtype=torch.double)
        basis_mat = self._get_lagrange_basis_mat(
            canonical_abscissa_1d, canonical_eval_samples)
        interp_vals = torch.mm(basis_mat, values)
        return interp_vals
//  [python]
def _cheby_interpolate(self, canonical_abscissa_1d,
                           canonical_barycentric_weights_1d, values,
                           canonical_eval_samples):
        # if type(values) != np.ndarray:
        #     values = values.detach().numpy()
        # interp_vals = multivariate_barycentric_lagrange_interpolation(
        #     canonical_eval_samples, canonical_abscissa_1d,
        #     canonical_barycentric_weights_1d, values,
        #     np.arange(self.nphys_vars))
        values = torch.as_tensor(values, dtype=torch.double)
        basis_mat = self._get_lagrange_basis_mat(
            canonical_abscissa_1d, canonical_eval_samples)
        interp_vals = torch.linalg.multi_dot((basis_mat, values))
        return interp_vals