jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def get_render_error_in_q(args, model, sample_size, device, targets, rgbs, poses, batch_size=1):
    ''' use nerf render imgs instead of use real imgs '''
    model.eval()
    
    results = np.zeros((sample_size, 2))
    print("to be implement...")

    predict_pose_list = []
    gt_pose_list = []
    ang_error_list = []

    for i in range(sample_size):
        data = rgbs[i:i+1].permute(0,3,1,2)
        pose = poses[i:i+1].reshape(batch_size, 12)

        data = data.to(device) # input
        pose = pose.reshape((batch_size,3,4)).numpy() # label

        # using SVD to make sure predict rotation is normalized rotation matrix
        with torch.no_grad():
            _, predict_pose = model(data)
            R_torch = predict_pose.reshape((batch_size, 3, 4))[:,:3,:3] # debug
            predict_pose = predict_pose.reshape((batch_size, 3, 4)).cpu().numpy()

            R = predict_pose[:,:3,:3]
            res = R@np.linalg.inv(R)
            # print('R@np.linalg.inv(R):', res)

            u,s,v=torch.svd(R_torch)
            Rs = torch.matmul(u, v.transpose(-2,-1))
        predict_pose[:,:3,:3] = Rs[:,:3,:3].cpu().numpy()

        pose_q = transforms.matrix_to_quaternion(torch.Tensor(pose[:,:3,:3]))#.cpu().numpy() # gnd truth in quaternion
        pose_x = pose[:, :3, 3] # gnd truth position
        predicted_q = transforms.matrix_to_quaternion(torch.Tensor(predict_pose[:,:3,:3]))#.cpu().numpy() # predict in quaternion
        predicted_x = predict_pose[:, :3, 3] # predict position
        pose_q = pose_q.squeeze() 
        pose_x = pose_x.squeeze() 
        predicted_q = predicted_q.squeeze() 
        predicted_x = predicted_x.squeeze()
        
        #Compute Individual Sample Error 
        q1 = pose_q / torch.linalg.norm(pose_q)
        q2 = predicted_q / torch.linalg.norm(predicted_q)
        d = torch.abs(torch.sum(torch.matmul(q1,q2))) 
        d = torch.clamp(d, -1., 1.) # acos can only input [-1~1]
        theta = (2 * torch.acos(d) * 180/math.pi).numpy()
        error_x = torch.linalg.norm(torch.Tensor(pose_x-predicted_x)).numpy()
        results[i,:] = [error_x, theta]
        #print ('Iteration: {} Error XYZ (m): {} Error Q (degrees): {}'.format(i, error_x, theta)) 

        # save results for visualization
        predict_pose_list.append(predicted_x)
        gt_pose_list.append(pose_x)
        ang_error_list.append(theta)

    predict_pose_list = np.array(predict_pose_list)
    gt_pose_list = np.array(gt_pose_list)
    ang_error_list = np.array(ang_error_list)
    vis_info_ret = {"pose": predict_pose_list, "pose_gt": gt_pose_list, "theta": ang_error_list}

    median_result = np.median(results,axis=0)
    mean_result = np.mean(results,axis=0)

    # standard log
    print ('Median error {}m and {} degrees.'.format(median_result[0], median_result[1]))
    print ('Mean error {}m and {} degrees.'.format(mean_result[0], mean_result[1]))

    return
//  [python]
def get_render_error_in_q(args, model, sample_size, device, targets, rgbs, poses, batch_size=1):
    ''' use nerf render imgs instead of use real imgs '''
    model.eval()
    
    results = np.zeros((sample_size, 2))
    print("to be implement...")

    predict_pose_list = []
    gt_pose_list = []
    ang_error_list = []

    for i in range(sample_size):
        data = rgbs[i:i+1].permute(0,3,1,2)
        pose = poses[i:i+1].reshape(batch_size, 12)

        data = data.to(device) # input
        pose = pose.reshape((batch_size,3,4)).numpy() # label

        # using SVD to make sure predict rotation is normalized rotation matrix
        with torch.no_grad():
            _, predict_pose = model(data)
            R_torch = predict_pose.reshape((batch_size, 3, 4))[:,:3,:3] # debug
            predict_pose = predict_pose.reshape((batch_size, 3, 4)).cpu().numpy()

            R = predict_pose[:,:3,:3]
            res = R@np.linalg.inv(R)
            # print('R@np.linalg.inv(R):', res)

            u,s,v=torch.linalg.svd(R_torch)
            Rs = torch.matmul(u, v.transpose(-2,-1))
        predict_pose[:,:3,:3] = Rs[:,:3,:3].cpu().numpy()

        pose_q = transforms.matrix_to_quaternion(torch.Tensor(pose[:,:3,:3]))#.cpu().numpy() # gnd truth in quaternion
        pose_x = pose[:, :3, 3] # gnd truth position
        predicted_q = transforms.matrix_to_quaternion(torch.Tensor(predict_pose[:,:3,:3]))#.cpu().numpy() # predict in quaternion
        predicted_x = predict_pose[:, :3, 3] # predict position
        pose_q = pose_q.squeeze() 
        pose_x = pose_x.squeeze() 
        predicted_q = predicted_q.squeeze() 
        predicted_x = predicted_x.squeeze()
        
        #Compute Individual Sample Error 
        q1 = pose_q / torch.linalg.norm(pose_q)
        q2 = predicted_q / torch.linalg.norm(predicted_q)
        d = torch.abs(torch.sum(torch.matmul(q1,q2))) 
        d = torch.clamp(d, -1., 1.) # acos can only input [-1~1]
        theta = (2 * torch.acos(d) * 180/math.pi).numpy()
        error_x = torch.linalg.norm(torch.Tensor(pose_x-predicted_x)).numpy()
        results[i,:] = [error_x, theta]
        #print ('Iteration: {} Error XYZ (m): {} Error Q (degrees): {}'.format(i, error_x, theta)) 

        # save results for visualization
        predict_pose_list.append(predicted_x)
        gt_pose_list.append(pose_x)
        ang_error_list.append(theta)

    predict_pose_list = np.array(predict_pose_list)
    gt_pose_list = np.array(gt_pose_list)
    ang_error_list = np.array(ang_error_list)
    vis_info_ret = {"pose": predict_pose_list, "pose_gt": gt_pose_list, "theta": ang_error_list}

    median_result = np.median(results,axis=0)
    mean_result = np.mean(results,axis=0)

    # standard log
    print ('Median error {}m and {} degrees.'.format(median_result[0], median_result[1]))
    print ('Mean error {}m and {} degrees.'.format(mean_result[0], mean_result[1]))

    return
//  [python]
def svd_on_cpu(self, matrix):
        u, s, v = torch.svd(matrix.to('cpu'))
        u = u.to(self.device)
        v = v.to(self.device)
        s = s.to(self.device)
        return u, s, v
//  [python]
def svd_on_cpu(self, matrix):
        u, s, v = torch.linalg.svd(matrix.to('cpu'))
        u = u.to(self.device)
        v = v.to(self.device)
        s = s.to(self.device)
        return u, s, v
//  [python]
def __compute_R(self,x) :
        Ji = torch.transpose(self.Ji_T,1,2)
        F = torch.matmul(Ji,x)
        u, s, v = torch.svd(F)
        R = torch.matmul(torch.transpose(u,1,2),v)
        return R
//  [python]
def __compute_R(self,x) :
        Ji = torch.transpose(self.Ji_T,1,2)
        F = torch.matmul(Ji,x)
        u, s, v = torch.linalg.svd(F)
        R = torch.matmul(torch.transpose(u,1,2),v)
        return R
//  [python]
def rigid_transform_3d(A, B, weights=None, weight_threshold=0):
    """
    Input:
        - A:       [bs, num_corr, 3], source point cloud
        - B:       [bs, num_corr, 3], target point cloud
        - weights: [bs, num_corr]     weight for each correspondence
        - weight_threshold: float,    clips points with weight below threshold
    Output:
        - R, t
    """
    bs = A.shape[0]
    if weights is None:
        weights = torch.ones_like(A[:, :, 0])
    weights[weights < weight_threshold] = 0
    # weights = weights / (torch.sum(weights, dim=-1, keepdim=True) + 1e-6)

    # find mean of point cloud
    centroid_A = torch.sum(A * weights[:, :, None], dim=1, keepdim=True) / (
            torch.sum(weights, dim=1, keepdim=True)[:, :, None] + 1e-6)
    centroid_B = torch.sum(B * weights[:, :, None], dim=1, keepdim=True) / (
            torch.sum(weights, dim=1, keepdim=True)[:, :, None] + 1e-6)

    # subtract mean
    Am = A - centroid_A
    Bm = B - centroid_B

    # construct weight covariance matrix
    Weight = torch.diag_embed(weights)  # 升维度，然后变为对角阵
    H = Am.permute(0, 2, 1) @ Weight @ Bm  # permute : tensor中的每一块做转置

    # find rotation
    U, S, Vt = torch.svd(H.cpu())
    U, S, Vt = U.to(weights.device), S.to(weights.device), Vt.to(weights.device)
    delta_UV = torch.det(Vt @ U.permute(0, 2, 1))
    eye = torch.eye(3)[None, :, :].repeat(bs, 1, 1).to(A.device)
    eye[:, -1, -1] = delta_UV
    R = Vt @ eye @ U.permute(0, 2, 1)
    t = centroid_B.permute(0, 2, 1) - R @ centroid_A.permute(0, 2, 1)
    # warp_A = transform(A, integrate_trans(R,t))
    # RMSE = torch.sum( (warp_A - B) ** 2, dim=-1).mean()
    return integrate_trans(R, t)
//  [python]
def rigid_transform_3d(A, B, weights=None, weight_threshold=0):
    """
    Input:
        - A:       [bs, num_corr, 3], source point cloud
        - B:       [bs, num_corr, 3], target point cloud
        - weights: [bs, num_corr]     weight for each correspondence
        - weight_threshold: float,    clips points with weight below threshold
    Output:
        - R, t
    """
    bs = A.shape[0]
    if weights is None:
        weights = torch.ones_like(A[:, :, 0])
    weights[weights < weight_threshold] = 0
    # weights = weights / (torch.sum(weights, dim=-1, keepdim=True) + 1e-6)

    # find mean of point cloud
    centroid_A = torch.sum(A * weights[:, :, None], dim=1, keepdim=True) / (
            torch.sum(weights, dim=1, keepdim=True)[:, :, None] + 1e-6)
    centroid_B = torch.sum(B * weights[:, :, None], dim=1, keepdim=True) / (
            torch.sum(weights, dim=1, keepdim=True)[:, :, None] + 1e-6)

    # subtract mean
    Am = A - centroid_A
    Bm = B - centroid_B

    # construct weight covariance matrix
    Weight = torch.diag_embed(weights)  # 升维度，然后变为对角阵
    H = Am.permute(0, 2, 1) @ Weight @ Bm  # permute : tensor中的每一块做转置

    # find rotation
    U, S, Vt = torch.linalg.svd(H.cpu())
    U, S, Vt = U.to(weights.device), S.to(weights.device), Vt.to(weights.device)
    delta_UV = torch.det(Vt @ U.permute(0, 2, 1))
    eye = torch.eye(3)[None, :, :].repeat(bs, 1, 1).to(A.device)
    eye[:, -1, -1] = delta_UV
    R = Vt @ eye @ U.permute(0, 2, 1)
    t = centroid_B.permute(0, 2, 1) - R @ centroid_A.permute(0, 2, 1)
    # warp_A = transform(A, integrate_trans(R,t))
    # RMSE = torch.sum( (warp_A - B) ** 2, dim=-1).mean()
    return integrate_trans(R, t)
//  [python]
def blurriness_svd(img, kr=10, sv_num=3):
    h, w = img.shape

    img = torch.nn.functional.pad(torch.from_numpy(img).float().cuda()[None, None], [kr] * 4, mode="reflect").squeeze()

    ys, xs = windowed_index(h, w, kr * 2)
    blocks = img[ys, xs]

    _, s, _ = torch.svd(blocks)

    top_svs = torch.sum(s[:, 0:sv_num], axis=1)
    total_svs = torch.sum(s, axis=1)
    sv_degrees = top_svs / total_svs
    max_sv = sv_degrees.min()
    min_sv = sv_degrees.max()

    blur_map = (sv_degrees - min_sv) / (max_sv - min_sv)

    return blur_map.float().reshape(h, w).cpu().numpy()
//  [python]
def blurriness_svd(img, kr=10, sv_num=3):
    h, w = img.shape

    img = torch.nn.functional.pad(torch.from_numpy(img).float().cuda()[None, None], [kr] * 4, mode="reflect").squeeze()

    ys, xs = windowed_index(h, w, kr * 2)
    blocks = img[ys, xs]

    _, s, _ = torch.linalg.svd(blocks)

    top_svs = torch.sum(s[:, 0:sv_num], axis=1)
    total_svs = torch.sum(s, axis=1)
    sv_degrees = top_svs / total_svs
    max_sv = sv_degrees.min()
    min_sv = sv_degrees.max()

    blur_map = (sv_degrees - min_sv) / (max_sv - min_sv)

    return blur_map.float().reshape(h, w).cpu().numpy()
//  [python]
def legalize_unitary(model: nn.Module):
    """
        Legalize the unitary matrices in the model.

        The function modifies the trainable unitary matrices in the model by applying
        singular value decomposition (SVD) and reassembling the matrices using the
        reconstructed singular values.

        Args:
            model (nn.Module): The model containing trainable unitary matrices.

        Returns:
            None

        Example:
            legalize_unitary(model)
        """
    with torch.no_grad():
        for name, params in model.named_parameters():
            if "TrainableUnitary" in name:
                U = params
                U, Sigma, V = torch.svd(U)
                params.data.copy_(U.matmul(V.conj().permute(0, 2, 1)))
//  [python]
def legalize_unitary(model: nn.Module):
    """
        Legalize the unitary matrices in the model.

        The function modifies the trainable unitary matrices in the model by applying
        singular value decomposition (SVD) and reassembling the matrices using the
        reconstructed singular values.

        Args:
            model (nn.Module): The model containing trainable unitary matrices.

        Returns:
            None

        Example:
            legalize_unitary(model)
        """
    with torch.no_grad():
        for name, params in model.named_parameters():
            if "TrainableUnitary" in name:
                U = params
                U, Sigma, V = torch.linalg.svd(U)
                params.data.copy_(U.matmul(V.conj().permute(0, 2, 1)))
//  [python]
def weighted_svd(self, pts1, pts2, weights, if_print=False):
        device = weights.device
        weights = weights.squeeze(1).unsqueeze(2)

        ones = torch.ones_like(weights)
        if self.is_cuda:
            ones = ones.cuda()
        pts1n, T1 = self.normalize(pts1, ones)
        pts2n, T2 = self.normalize(pts2, ones)
        # pts1n, T1 = self.normalize(pts1, weights)
        # pts2n, T2 = self.normalize(pts2, weights)


        p = torch.cat((pts2n[:,0].unsqueeze(1)*pts1n,
                       pts2n[:,1].unsqueeze(1)*pts1n,
                       pts1n), 1).permute(0,2,1)

        # # if self.normalize_SVD:
        # #     p = torch.nn.functional.normalize(p, dim=2)
        # X = p*torch.sqrt(weights)

        if self.normalize_SVD:
            p = torch.nn.functional.normalize(p, dim=2)
        X = p*weights

        out_b = []
        F_vecs_list = []
        if self.if_cpu_svd:
            for b in range(X.size(0)):
                _, _, V = torch.svd(X[b].cpu())
                F = V[:,-1].view(3,3)
                F_vecs_list.append(V[:,-1]/(V[:,-1].norm()))
                U, S, V = torch.svd(F)
                F_ = U.mm((S*self.mask.cpu()).diag()).mm(V.t())
                out_b.append(F_.unsqueeze(0))
            out = torch.cat(out_b, 0).cuda()
            F_vecs= torch.stack(F_vecs_list).cuda()
        else:
            for b in range(X.size(0)):
                _, _, V = torch.svd(X[b])
                F = V[:,-1].view(3,3)
                F_vecs_list.append(V[:,-1]/(V[:,-1].norm()))
                U, S, V = torch.svd(F)
                F_ = U.mm((S*self.mask.to(device)).diag()).mm(V.t())
                out_b.append(F_.unsqueeze(0))
            out = torch.cat(out_b, 0)
            F_vecs = torch.stack(F_vecs_list)

        if if_print:
            print(F_vecs.size(), p.size(), weights.size())
            print('----F_vecs')
            print(F_vecs[0].detach().cpu().numpy())
            print('----p')
            print(p[0].detach().cpu().numpy())
            print('----weights')
            print(weights[:2].squeeze().detach().cpu().numpy(), torch.sum(weights[:2], dim=1).squeeze().detach().cpu().numpy())

        residual = (X @ F_vecs.unsqueeze(-1)).squeeze(-1) # [B, N, 1]
        # residual_nonWeighted = (p @ F_vecs.unsqueeze(-1)).squeeze(-1) # [B, N, 1]
        # print(residual.size())
        # print(residual.norm(p=2, dim=1).size())

        out = T2.permute(0,2,1).bmm(out).bmm(T1)

        return out, residual.squeeze(-1)
//  [python]
def weighted_svd(self, pts1, pts2, weights, if_print=False):
        device = weights.device
        weights = weights.squeeze(1).unsqueeze(2)

        ones = torch.ones_like(weights)
        if self.is_cuda:
            ones = ones.cuda()
        pts1n, T1 = self.normalize(pts1, ones)
        pts2n, T2 = self.normalize(pts2, ones)
        # pts1n, T1 = self.normalize(pts1, weights)
        # pts2n, T2 = self.normalize(pts2, weights)


        p = torch.cat((pts2n[:,0].unsqueeze(1)*pts1n,
                       pts2n[:,1].unsqueeze(1)*pts1n,
                       pts1n), 1).permute(0,2,1)

        # # if self.normalize_SVD:
        # #     p = torch.nn.functional.normalize(p, dim=2)
        # X = p*torch.sqrt(weights)

        if self.normalize_SVD:
            p = torch.nn.functional.normalize(p, dim=2)
        X = p*weights

        out_b = []
        F_vecs_list = []
        if self.if_cpu_svd:
            for b in range(X.size(0)):
                U, S, V = torch.linalg.svd(X[b].cpu())
                F = V[:,-1].view(3,3)
                F_vecs_list.append(V[:,-1]/(V[:,-1].norm()))
                U, S, V = torch.linalg.svd(F)
                F_ = U.mm((S*self.mask.cpu()).diag()).mm(V.t())
                out_b.append(F_.unsqueeze(0))
            out = torch.cat(out_b, 0).cuda()
            F_vecs= torch.stack(F_vecs_list).cuda()
        else:
            for b in range(X.size(0)):
                U, S, V = torch.linalg.svd(X[b])
                F = V[:,-1].view(3,3)
                F_vecs_list.append(V[:,-1]/(V[:,-1].norm()))
                U, S, V = torch.linalg.svd(F)
                F_ = U.mm((S*self.mask.to(device)).diag()).mm(V.t())
                out_b.append(F_.unsqueeze(0))
            out = torch.cat(out_b, 0)
            F_vecs = torch.stack(F_vecs_list)

        if if_print:
            print(F_vecs.size(), p.size(), weights.size())
            print('----F_vecs')
            print(F_vecs[0].detach().cpu().numpy())
            print('----p')
            print(p[0].detach().cpu().numpy())
            print('----weights')
            print(weights[:2].squeeze().detach().cpu().numpy(), torch.sum(weights[:2], dim=1).squeeze().detach().cpu().numpy())

        residual = (X @ F_vecs.unsqueeze(-1)).squeeze(-1) # [B, N, 1]
        # residual_nonWeighted = (p @ F_vecs.unsqueeze(-1)).squeeze(-1) # [B, N, 1]
        # print(residual.size())
        # print(residual.norm(p=2, dim=1).size())

        out = T2.permute(0,2,1).bmm(out).bmm(T1)

        return out, residual.squeeze(-1)
//  [python]
def reduce(self, grad_in, grad_out, memory_out):
        """
        Reduce gradients between the workers in place
        :param grad_in: dictionary
        :param grad_out: dictionary
        :param memory_out: dictionary
        """
        bits_communicated = 0

        # Split the tensors into rank1-ones that will be reduced un-compressed
        # and rank > 1 tensors that are compressed
        rank1_tensors = [
            (tensor, out, mem)
            for tensor, out, mem in zip(grad_in, grad_out, memory_out)
            if tensor.ndimension() <= 1
        ]
        high_rank_tensors = [
            (tensor, out, mem)
            for tensor, out, mem in zip(grad_in, grad_out, memory_out)
            if tensor.ndimension() > 1
        ]

        # Communicate rank 1 tensors
        bits_communicated += self._reduce_rank1(rank1_tensors)

        for tensor, out, mem in high_rank_tensors:
            m = tensor.shape[0]
            n = tensor.nelement() // m

            matrix = tensor.view(tensor.shape[0], -1)

            rnk = min(self.rank, m, n)
            u, s, v = torch.svd(matrix)
            u, s, v = u[:, :rnk], s[:rnk], v[:, :rnk]

            mem.data[:] = tensor
            mem.view(*matrix.shape).data -= torch.einsum("in, n, jn -> ij", u, s, v)

            if self.n_workers > 1:
                worker_u = [torch.empty_like(u) for i in range(self.n_workers)]
                worker_v = [torch.empty_like(v) for i in range(self.n_workers)]
                worker_s = [torch.empty_like(s) for i in range(self.n_workers)]
                h1 = all_gather(worker_u, u, async_op=True)
                h2 = all_gather(worker_v, v, async_op=True)
                h3 = all_gather(worker_s, s, async_op=True)
                h1.wait()
                h2.wait()
                h3.wait()
            else:
                worker_u = [u]
                worker_v = [v]
                worker_s = [s]

            out.data[:] = 0.0
            for uu, ss, vv in zip(worker_u, worker_s, worker_v):
                out.view(*matrix.shape).add_(
                    1.0 / self.n_workers, torch.einsum("in, n, jn -> ij", uu, ss, vv)
                )

            bits_communicated += n_bits(u) + n_bits(s) + n_bits(v)

        return bits_communicated
//  [python]
def reduce(self, grad_in, grad_out, memory_out):
        """
        Reduce gradients between the workers in place
        :param grad_in: dictionary
        :param grad_out: dictionary
        :param memory_out: dictionary
        """
        bits_communicated = 0

        # Split the tensors into rank1-ones that will be reduced un-compressed
        # and rank > 1 tensors that are compressed
        rank1_tensors = [
            (tensor, out, mem)
            for tensor, out, mem in zip(grad_in, grad_out, memory_out)
            if tensor.ndimension() <= 1
        ]
        high_rank_tensors = [
            (tensor, out, mem)
            for tensor, out, mem in zip(grad_in, grad_out, memory_out)
            if tensor.ndimension() > 1
        ]

        # Communicate rank 1 tensors
        bits_communicated += self._reduce_rank1(rank1_tensors)

        for tensor, out, mem in high_rank_tensors:
            m = tensor.shape[0]
            n = tensor.nelement() // m

            matrix = tensor.view(tensor.shape[0], -1)

            rnk = min(self.rank, m, n)
            u, s, v = torch.linalg.svd(matrix)
            u, s, v = u[:, :rnk], s[:rnk], v[:, :rnk]

            mem.data[:] = tensor
            mem.view(*matrix.shape).data -= torch.einsum("in, n, jn -> ij", u, s, v)

            if self.n_workers > 1:
                worker_u = [torch.empty_like(u) for i in range(self.n_workers)]
                worker_v = [torch.empty_like(v) for i in range(self.n_workers)]
                worker_s = [torch.empty_like(s) for i in range(self.n_workers)]
                h1 = all_gather(worker_u, u, async_op=True)
                h2 = all_gather(worker_v, v, async_op=True)
                h3 = all_gather(worker_s, s, async_op=True)
                h1.wait()
                h2.wait()
                h3.wait()
            else:
                worker_u = [u]
                worker_v = [v]
                worker_s = [s]

            out.data[:] = 0.0
            for uu, ss, vv in zip(worker_u, worker_s, worker_v):
                out.view(*matrix.shape).add_(
                    1.0 / self.n_workers, torch.einsum("in, n, jn -> ij", uu, ss, vv)
                )

            bits_communicated += n_bits(u) + n_bits(s) + n_bits(v)

        return bits_communicated
//  [python]
def batch_get_children_orient_svd(rel_pose_skeleton, rel_rest_pose, rot_mat_chain_parent, children_list, dtype):
    # if isinstance(rel_pose_skeleton, list):
    #     device = rel_pose_skeleton[0].device
    # else:
    #     device = rel_pose_skeleton.device

    rest_mat = []
    target_mat = []
    for c, child in enumerate(children_list):
        if isinstance(rel_pose_skeleton, list):
            target = rel_pose_skeleton[c].clone()
            template = rel_rest_pose[c].clone()
        else:
            target = rel_pose_skeleton[:, child].clone()
            template = rel_rest_pose[:, child].clone()

        target = torch.matmul(
            rot_mat_chain_parent.transpose(1, 2),
            target)

        target_mat.append(target)
        rest_mat.append(template)

    rest_mat = torch.cat(rest_mat, dim=2)
    target_mat = torch.cat(target_mat, dim=2)
    S = rest_mat.bmm(target_mat.transpose(1, 2))

    U, Sig, V = torch.svd(S)
    # U, Sig, V = torch.svd(S.cpu())
    # U = U.to(device=device)
    # V = V.to(device=device)

    # assert (Sig[:, 0] >= Sig[:, 1]).all() and (Sig[:, 1] >= Sig[:, 2]).all(), Sig

    det_u_vt = torch.det(torch.bmm(V, U.transpose(1, 2)))
    det_modify_mat = torch.eye(3, device=U.device).unsqueeze(0).expand(U.shape[0], -1, -1).clone()
    det_modify_mat[:, 2, 2] = det_u_vt
    rot_mat = torch.bmm(torch.bmm(V, det_modify_mat), U.transpose(1, 2))
    assert torch.sum(torch.isnan(rot_mat)) == 0, ('3children rot_mat', rot_mat)

    return rot_mat
//  [python]
def batch_get_children_orient_svd(rel_pose_skeleton, rel_rest_pose, rot_mat_chain_parent, children_list, dtype):
    rest_mat = []
    target_mat = []
    for c, child in enumerate(children_list):
        if isinstance(rel_pose_skeleton, list):
            target = rel_pose_skeleton[c].clone()
            template = rel_rest_pose[c].clone()
        else:
            target = rel_pose_skeleton[:, child].clone()
            template = rel_rest_pose[:, child].clone()

        target = torch.matmul(
            rot_mat_chain_parent.transpose(1, 2),
            target)

        target_mat.append(target)
        rest_mat.append(template)

    rest_mat = torch.cat(rest_mat, dim=2)
    target_mat = torch.cat(target_mat, dim=2)
    S = rest_mat.bmm(target_mat.transpose(1, 2))

    U, S, V = torch.linalg.svd(S)

    det_u_vt = torch.det(torch.bmm(V, U.transpose(1, 2)))
    det_modify_mat = torch.eye(3, device=U.device).unsqueeze(0).expand(U.shape[0], -1, -1).clone()
    det_modify_mat[:, 2, 2] = det_u_vt
    rot_mat = torch.bmm(torch.bmm(V, det_modify_mat), U.transpose(1, 2))
    assert torch.sum(torch.isnan(rot_mat)) == 0, ('3children rot_mat', rot_mat)

    return rot_mat
//  [python]
def batch_get_pelvis_orient_svd(rel_pose_skeleton, rel_rest_pose, parents, children, dtype):
    # if isinstance(rel_pose_skeleton, list):
    #     device = rel_pose_skeleton[0].device
    # else:
    #     device = rel_pose_skeleton.device

    pelvis_child = []
    for i in range(parents.shape[0]):
        if parents[i] == 0 and i not in pelvis_child:
            pelvis_child.append(i)

    rest_mat = []
    target_mat = []
    for child in pelvis_child:
        rest_mat.append(rel_rest_pose[:, child].clone())
        target_mat.append(rel_pose_skeleton[:, child].clone())

    rest_mat = torch.cat(rest_mat, dim=2)
    target_mat = torch.cat(target_mat, dim=2)
    S = rest_mat.bmm(target_mat.transpose(1, 2))

    mask_zero = S.sum(dim=(1, 2))

    S_non_zero = S[mask_zero != 0].reshape(-1, 3, 3)

    U, Sig, V = torch.svd(S_non_zero)
    # U, Sig, V = torch.svd(S_non_zero.cpu())
    # U = U.to(device=device)
    # V = V.to(device=device)

    assert (Sig[:, 0] >= Sig[:, 1]).all() and (Sig[:, 1] >= Sig[:, 2]).all(), Sig

    rot_mat = torch.zeros_like(S)
    rot_mat[mask_zero == 0] = torch.eye(3, device=S.device)

    # rot_mat_non_zero = torch.bmm(V, U.transpose(1, 2))

    det_u_vt = torch.det(torch.bmm(V, U.transpose(1, 2)))
    det_modify_mat = torch.eye(3, device=U.device).unsqueeze(0).expand(U.shape[0], -1, -1).clone()
    det_modify_mat[:, 2, 2] = det_u_vt
    rot_mat_non_zero = torch.bmm(torch.bmm(V, det_modify_mat), U.transpose(1, 2))

    rot_mat[mask_zero != 0] = rot_mat_non_zero

    assert torch.sum(torch.isnan(rot_mat)) == 0, ('rot_mat', rot_mat)

    return rot_mat
//  [python]
def batch_get_pelvis_orient_svd(rel_pose_skeleton, rel_rest_pose, parents, children, dtype):
    pelvis_child = []
    for i in range(parents.shape[0]):
        if parents[i] == 0 and i not in pelvis_child:
            pelvis_child.append(i)

    rest_mat = []
    target_mat = []
    for child in pelvis_child:
        rest_mat.append(rel_rest_pose[:, child].clone())
        target_mat.append(rel_pose_skeleton[:, child].clone())

    rest_mat = torch.cat(rest_mat, dim=2)
    target_mat = torch.cat(rest_mat, dim=2)
    S = rest_mat.bmm(target_mat.transpose(1, 2))

    mask_zero = S.sum(dim=(1, 2))

    S_non_zero = S[mask_zero != 0].reshape(-1, 3, 3)

    U, S, V = torch.linalg.svd(S_non_zero)

    assert (S[:, 0] >= S[:, 1]).all() and (S[:, 1] >= S[:, 2]).all(), S

    rot_mat = torch.zeros_like(S)
    rot_mat[mask_zero == 0] = torch.eye(3, device=S.device)

    det_u_vt = torch.det(torch.bmm(V, U.transpose(1, 2)))
    det_modify_mat = torch.eye(3, device=U.device).unsqueeze(0).expand(U.shape[0], -1, -1).clone()
    det_modify_mat[:, 2, 2] = det_u_vt
    rot_mat_non_zero = torch.bmm(torch.bmm(V, det_modify_mat), U.transpose(1, 2))

    rot_mat[mask_zero != 0] = rot_mat_non_zero

    assert torch.sum(torch.isnan(rot_mat)) == 0, ('rot_mat', rot_mat)

    return rot_mat
//  [python]
def _loss(self, train_batch, batch_results, batch_idx):
        #### Prepare ####
        system = self.get_system()

        ## Batch
        batch = self.get_batch(train_batch, batch_idx)
        rays = batch['coords']

        ## tform constraints
        raw = system.render(
            'embed_params',
            rays
        )['value']
        out_channels = (raw.shape[-1] // (self.param_channels + 1))
        tform = raw[..., :out_channels].reshape(
            -1, out_channels, self.param_channels
        )

        _, S, _ = torch.svd(tform)

        loss = self.loss_fn(
            S[..., 2:],
            torch.zeros_like(S[..., 2:])
        )

        return loss
//  [python]
def _loss(self, train_batch, batch_results, batch_idx):
        #### Prepare ####
        system = self.get_system()

        ## Batch
        batch = self.get_batch(train_batch, batch_idx)
        rays = batch['coords']

        ## tform constraints
        raw = system.render(
            'embed_params',
            rays
        )['value']
        out_channels = (raw.shape[-1] // (self.param_channels + 1))
        tform = raw[..., :out_channels].reshape(
            -1, out_channels, self.param_channels
        )

        _, S, _ = torch.linalg.svd(tform)

        loss = self.loss_fn(
            S[..., 2:],
            torch.zeros_like(S[..., 2:])
        )

        return loss
//  [python]
def nearsub(train_features, train_labels, test_features, test_labels, 
            num_classes, n_comp=10, return_pred=False):
    train_scores, test_scores = [], []
    classes = np.arange(num_classes)
    features_sort, _ = utils.sort_dataset(train_features, train_labels, 
                                          classes=classes, stack=False)           
    fd = features_sort[0].shape[1]
    for j in classes:
        _, _, V = torch.svd(features_sort[j])
        components = V[:, :n_comp].T
        subspace_j = torch.eye(fd) - components.T @ components
        train_j = subspace_j @ train_features.T
        test_j = subspace_j @ test_features.T
        train_scores_j = torch.linalg.norm(train_j, ord=2, axis=0)
        test_scores_j = torch.linalg.norm(test_j, ord=2, axis=0)
        train_scores.append(train_scores_j)
        test_scores.append(test_scores_j)
    train_pred = torch.stack(train_scores).argmin(0)
    test_pred = torch.stack(test_scores).argmin(0)
    if return_pred:
        return train_pred.numpy(), test_pred.numpy()
    train_acc = compute_accuracy(classes[train_pred], train_labels.numpy())
    test_acc = compute_accuracy(classes[test_pred], test_labels.numpy())
    print('SVD: {}, {}'.format(train_acc, test_acc))
    return train_acc, test_acc
//  [python]
def nearsub(train_features, train_labels, test_features, test_labels, 
            num_classes, n_comp=10, return_pred=False):
    train_scores, test_scores = [], []
    classes = np.arange(num_classes)
    features_sort, _ = utils.sort_dataset(train_features, train_labels, 
                                          classes=classes, stack=False)           
    fd = features_sort[0].shape[1]
    for j in classes:
        U, S, V = torch.linalg.svd(features_sort[j])
        components = V[:, :n_comp].T
        subspace_j = torch.eye(fd) - components.T @ components
        train_j = subspace_j @ train_features.T
        test_j = subspace_j @ test_features.T
        train_scores_j = torch.linalg.norm(train_j, ord=2, axis=0)
        test_scores_j = torch.linalg.norm(test_j, ord=2, axis=0)
        train_scores.append(train_scores_j)
        test_scores.append(test_scores_j)
    train_pred = torch.stack(train_scores).argmin(0)
    test_pred = torch.stack(test_scores).argmin(0)
    if return_pred:
        return train_pred.numpy(), test_pred.numpy()
    train_acc = compute_accuracy(classes[train_pred], train_labels.numpy())
    test_acc = compute_accuracy(classes[test_pred], test_labels.numpy())
    print('SVD: {}, {}'.format(train_acc, test_acc))
    return train_acc, test_acc
//  [python]
def __init__(self, A):
        self._U, self._singulars, self._V = torch.svd(A, some=False)
        self._Vt = self._V.transpose(0, 1)
        self._Ut = self._U.transpose(0, 1)

        ZERO = 1e-3
        self._singulars[self._singulars < ZERO] = 0
        print(len([x.item() for x in self._singulars if x == 0]))
//  [python]
def __init__(self, A):
        self._U, self._singulars, self._V = torch.linalg.svd(A, full_matrices=False)
        self._Vt = self._V.transpose(0, 1)
        self._Ut = self._U.transpose(0, 1)

        ZERO = 1e-3
        self._singulars[self._singulars < ZERO] = 0
        print(len([x.item() for x in self._singulars if x == 0]))
//  [python]
def __init__(self, channels, img_dim, ratio, device): #ratio = 2 or 4
        self.img_dim = img_dim
        self.channels = channels
        self.y_dim = img_dim // 32
        self.ratio = 32
        A = torch.randn(32**2, 32**2).to(device)
        _, _, self.V_small = torch.svd(A, some=False)
        self.Vt_small = self.V_small.transpose(0, 1)
        self.singulars_small = torch.ones(int(32 * 32 * ratio), device=device)
        self.cs_size = self.singulars_small.size(0)
//  [python]
def __init__(self, channels, img_dim, ratio, device): #ratio = 2 or 4
        self.img_dim = img_dim
        self.channels = channels
        self.y_dim = img_dim // 32
        self.ratio = 32
        A = torch.randn(32**2, 32**2).to(device)
        _, self.V_small = torch.linalg.svd(A, full_matrices=False)
        self.Vt_small = self.V_small.transpose(0, 1)
        self.singulars_small = torch.ones(int(32 * 32 * ratio), device=device)
        self.cs_size = self.singulars_small.size(0)
//  [python]
def diffusion_distance(mat, sigma=8.0, alpha=1.0):
    D =distance_matrix(mat);
    K = torch.exp(-(torch.pow(torch.div(D,sigma) ,2))) # Kernel
    p = K.sum(1)
    K1 = K/(torch.pow(p.unsqueeze(1)*p,alpha)+1e-9) # alpha = 1 Laplace Beltrami, 0.5 Fokker Planck diffusion.
    v = torch.sqrt(K1.sum(1))
    A = K1/(1e-9+v.unsqueeze(1)*v)
    [u,s,v]=torch.svd(A)
    u=u/(1e-9+u[:,0].unsqueeze(1))
    return K1,u,s
//  [python]
def diffusion_distance(mat, sigma=8.0, alpha=1.0):
    D =distance_matrix(mat);
    K = torch.exp(-(torch.pow(torch.div(D,sigma) ,2))) # Kernel
    p = K.sum(1)
    K1 = K/(torch.pow(p.unsqueeze(1)*p,alpha)+1e-9) # alpha = 1 Laplace Beltrami, 0.5 Fokker Planck diffusion.
    v = torch.sqrt(K1.sum(1))
    A = K1/(1e-9+v.unsqueeze(1)*v)
    [u,s,v]=torch.linalg.svd(A)
    u=u/(1e-9+u[:,0].unsqueeze(1))
    return K1,u,s
//  [python]
def __init__(self, channels, img_dim, ratio, device): #ratio = 2 or 4
        assert img_dim % ratio == 0
        self.img_dim = img_dim
        self.channels = channels
        self.y_dim = img_dim // ratio
        self.ratio = ratio
        A = torch.Tensor([[1 / ratio**2] * ratio**2]).to(device)
        self.U_small, self.singulars_small, self.V_small = torch.svd(A, some=False)
        self.Vt_small = self.V_small.transpose(0, 1)
//  [python]
def __init__(self, channels, img_dim, ratio, device): #ratio = 2 or 4
        assert img_dim % ratio == 0
        self.img_dim = img_dim
        self.channels = channels
        self.y_dim = img_dim // ratio
        self.ratio = ratio
        A = torch.Tensor([[1 / ratio**2] * ratio**2]).to(device)
        self.U_small, self.singulars_small, self.V_small = torch.linalg.svd(A)
        self.Vt_small = self.V_small.transpose(0, 1)
//  [python]
def __init__(self, img_dim, device):
        self.channels = 3
        self.img_dim = img_dim
        #Do the SVD for the per-pixel matrix
        A = torch.Tensor([[0.3333, 0.3334, 0.3333]]).to(device)
        self.U_small, self.singulars_small, self.V_small = torch.svd(A, some=False)
        self.Vt_small = self.V_small.transpose(0, 1)
//  [python]
def __init__(self, img_dim, device):
        self.channels = 3
        self.img_dim = img_dim
        #Do the SVD for the per-pixel matrix
        A = torch.Tensor([[0.3333, 0.3334, 0.3333]]).to(device)
        self.U_small, self.singulars_small, self.V_small = torch.linalg.svd(A)
        self.Vt_small = self.V_small.transpose(0, 1)
//  [python]
def tensor_log(self, t):#4dim
        u, s, v = torch.svd(t)
        return u @ torch.diag_embed(torch.log(s)) @ v.permute(0, 1, 3, 2)
//  [python]
def tensor_log(self, t):#4dim
        u, s, v = torch.linalg.svd(t)
        return u @ torch.diag_embed(torch.log(s)) @ v.permute(0, 1, 3, 2)
//  [python]
def log_euclidean_distance(self, A, B):
        inner_term = self.tensor_log(A) - self.tensor_log(B)
        inner_multi = inner_term @ inner_term.permute(0, 1, 3, 2)
        _, s, _= torch.svd(inner_multi)
        final = torch.sum(s, dim=-1)
        return final
//  [python]
def log_euclidean_distance(self, A, B):
        inner_term = self.tensor_log(A) - self.tensor_log(B)
        inner_multi = inner_term @ inner_term.permute(0, 1, 3, 2)
        _, s, _= torch.linalg.svd(inner_multi)
        final = torch.sum(s, dim=-1)
        return final
//  [python]
def __init__(self, kernel, channels, img_dim, device, stride = 1):
        self.img_dim = img_dim
        self.channels = channels
        self.ratio = stride
        small_dim = img_dim // stride
        self.small_dim = small_dim
        #build 1D conv matrix
        A_small = torch.zeros(small_dim, img_dim, device=device)
        for i in range(stride//2, img_dim + stride//2, stride):
            for j in range(i - kernel.shape[0]//2, i + kernel.shape[0]//2):
                j_effective = j
                #reflective padding
                if j_effective < 0: j_effective = -j_effective-1
                if j_effective >= img_dim: j_effective = (img_dim - 1) - (j_effective - img_dim)
                #matrix building
                A_small[i // stride, j_effective] += kernel[j - i + kernel.shape[0]//2]
        #get the svd of the 1D conv
        self.U_small, self.singulars_small, self.V_small = torch.svd(A_small, some=False)
        ZERO = 3e-2
        self.singulars_small[self.singulars_small < ZERO] = 0
        #calculate the singular values of the big matrix
        self._singulars = torch.matmul(self.singulars_small.reshape(small_dim, 1), self.singulars_small.reshape(1, small_dim)).reshape(small_dim**2)
        #permutation for matching the singular values. See P_1 in Appendix D.5.
        self._perm = torch.Tensor([self.img_dim * i + j for i in range(self.small_dim) for j in range(self.small_dim)] + \
                                  [self.img_dim * i + j for i in range(self.small_dim) for j in range(self.small_dim, self.img_dim)]).to(device).long()
//  [python]
def __init__(self, kernel, channels, img_dim, device, stride = 1):
        self.img_dim = img_dim
        self.channels = channels
        self.ratio = stride
        small_dim = img_dim // stride
        self.small_dim = small_dim
        #build 1D conv matrix
        A_small = torch.zeros(small_dim, img_dim, device=device)
        for i in range(stride//2, img_dim + stride//2, stride):
            for j in range(i - kernel.shape[0]//2, i + kernel.shape[0]//2):
                j_effective = j
                #reflective padding
                if j_effective < 0: j_effective = -j_effective-1
                if j_effective >= img_dim: j_effective = (img_dim - 1) - (j_effective - img_dim)
                #matrix building
                A_small[i // stride, j_effective] += kernel[j - i + kernel.shape[0]//2]
        #get the svd of the 1D conv
        self.U_small, self.singulars_small, self.V_small = torch.linalg.svd(A_small)
        ZERO = 3e-2
        self.singulars_small[self.singulars_small < ZERO] = 0
        #calculate the singular values of the big matrix
        self._singulars = torch.matmul(self.singulars_small.reshape(small_dim, 1), self.singulars_small.reshape(1, small_dim)).reshape(small_dim**2)
        #permutation for matching the singular values. See P_1 in Appendix D.5.
        self._perm = torch.Tensor([self.img_dim * i + j for i in range(self.small_dim) for j in range(self.small_dim)] + \
                                  [self.img_dim * i + j for i in range(self.small_dim) for j in range(self.small_dim, self.img_dim)]).to(device).long()
//  [python]
def __init__(self, kernel, channels, img_dim, device, ZERO = 3e-2):
        self.img_dim = img_dim
        self.channels = channels
        #build 1D conv matrix
        A_small = torch.zeros(img_dim, img_dim, device=device)
        for i in range(img_dim):
            for j in range(i - kernel.shape[0]//2, i + kernel.shape[0]//2):
                if j < 0 or j >= img_dim: continue
                A_small[i, j] = kernel[j - i + kernel.shape[0]//2]
        #get the svd of the 1D conv
        self.U_small, self.singulars_small, self.V_small = torch.svd(A_small, some=False)
        #ZERO = 3e-2
        self.singulars_small_orig = self.singulars_small.clone()
        self.singulars_small[self.singulars_small < ZERO] = 0
        #calculate the singular values of the big matrix
        self._singulars_orig = torch.matmul(self.singulars_small_orig.reshape(img_dim, 1), self.singulars_small_orig.reshape(1, img_dim)).reshape(img_dim**2)
        self._singulars = torch.matmul(self.singulars_small.reshape(img_dim, 1), self.singulars_small.reshape(1, img_dim)).reshape(img_dim**2)
        #sort the big matrix singulars and save the permutation
        self._singulars, self._perm = self._singulars.sort(descending=True) #, stable=True)
        self._singulars_orig = self._singulars_orig[self._perm]
//  [python]
def __init__(self, kernel, channels, img_dim, device, ZERO = 3e-2):
        self.img_dim = img_dim
        self.channels = channels
        #build 1D conv matrix
        A_small = torch.zeros(img_dim, img_dim, device=device)
        for i in range(img_dim):
            for j in range(i - kernel.shape[0]//2, i + kernel.shape[0]//2):
                if j < 0 or j >= img_dim: continue
                A_small[i, j] = kernel[j - i + kernel.shape[0]//2]
        #get the svd of the 1D conv
        self.U_small, self.singulars_small, self.V_small = torch.linalg.svd(A_small)
        #ZERO = 3e-2
        self.singulars_small_orig = self.singulars_small.clone()
        self.singulars_small[self.singulars_small < ZERO] = 0
        #calculate the singular values of the big matrix
        self._singulars_orig = torch.matmul(self.singulars_small_orig.reshape(img_dim, 1), self.singulars_small_orig.reshape(1, img_dim)).reshape(img_dim**2)
        self._singulars = torch.matmul(self.singulars_small.reshape(img_dim, 1), self.singulars_small.reshape(1, img_dim)).reshape(img_dim**2)
        #sort the big matrix singulars and save the permutation
        self._singulars, self._perm = self._singulars.sort(descending=True) #, stable=True)
        self._singulars_orig = self._singulars_orig[self._perm]
//  [python]
def compute_focal_lengths(fund, Issize, Itsize):
    F_tensor = torch.tensor(fund).float()
    
    f_t = torch.nn.Parameter(max(Issize) * torch.ones(1))
    f_s = torch.nn.Parameter(max(Itsize) * torch.ones(1))
    f_s.require_grad = True
    f_t.require_grad = True

    opt = torch.optim.SGD([f_s, f_t], lr=100, momentum=0.9)

    for lr in [1, .1]:
        opt.lr = lr
        for i in range(2000):
            K1 = torch.eye(3)
            K1[0, 0] = f_s
            K1[1, 1] = f_s

            K2 = torch.eye(3)
            K2[0, 0] = f_t
            K2[1, 1] = f_t

            E = K2.t() @ F_tensor @ K1
            u, s, v = torch.svd(E / torch.norm(E))
            l = s[0] - s[1]
            opt.zero_grad()
            l.backward()
            opt.step()

    return f_s.item(), f_t.item()
//  [python]
def compute_focal_lengths(fund, Issize, Itsize):
    F_tensor = torch.tensor(fund).float()
    
    f_t = torch.nn.Parameter(max(Issize) * torch.ones(1))
    f_s = torch.nn.Parameter(max(Itsize) * torch.ones(1))
    f_s.require_grad = True
    f_t.require_grad = True

    opt = torch.optim.SGD([f_s, f_t], lr=100, momentum=0.9)

    for lr in [1, .1]:
        opt.lr = lr
        for i in range(2000):
            K1 = torch.eye(3)
            K1[0, 0] = f_s
            K1[1, 1] = f_s

            K2 = torch.eye(3)
            K2[0, 0] = f_t
            K2[1, 1] = f_t

            E = K2.t() @ F_tensor @ K1
            u, s, v = torch.linalg.svd(E / torch.norm(E))
            l = s[0] - s[1]
            opt.zero_grad()
            l.backward()
            opt.step()

    return f_s.item(), f_t.item()
//  [python]
def __init__(self, kernel1, kernel2, channels, img_dim, device):
        self.img_dim = img_dim
        self.channels = channels
        A_small1 = torch.zeros(img_dim, img_dim, device=device)
        for i in range(img_dim):
            for j in range(i - kernel1.shape[0]//2, i + kernel1.shape[0]//2):
                if j < 0 or j >= img_dim: continue
                A_small1[i, j] = kernel1[j - i + kernel1.shape[0]//2]
        A_small2 = torch.zeros(img_dim, img_dim, device=device)
        for i in range(img_dim):
            for j in range(i - kernel2.shape[0]//2, i + kernel2.shape[0]//2):
                if j < 0 or j >= img_dim: continue
                A_small2[i, j] = kernel2[j - i + kernel2.shape[0]//2]
        self.U_small1, self.singulars_small1, self.V_small1 = torch.svd(A_small1, some=False)
        self.U_small2, self.singulars_small2, self.V_small2 = torch.svd(A_small2, some=False)
        ZERO = 3e-2
        self.singulars_small1[self.singulars_small1 < ZERO] = 0
        self.singulars_small2[self.singulars_small2 < ZERO] = 0
        #calculate the singular values of the big matrix
        self._singulars = torch.matmul(self.singulars_small1.reshape(img_dim, 1), self.singulars_small2.reshape(1, img_dim)).reshape(img_dim**2)
        #sort the big matrix singulars and save the permutation
        self._singulars, self._perm = self._singulars.sort(descending=True)
//  [python]
def __init__(self, kernel1, kernel2, channels, img_dim, device):
        self.img_dim = img_dim
        self.channels = channels
        A_small1 = torch.zeros(img_dim, img_dim, device=device)
        for i in range(img_dim):
            for j in range(i - kernel1.shape[0]//2, i + kernel1.shape[0]//2):
                if j < 0 or j >= img_dim: continue
                A_small1[i, j] = kernel1[j - i + kernel1.shape[0]//2]
        A_small2 = torch.zeros(img_dim, img_dim, device=device)
        for i in range(img_dim):
            for j in range(i - kernel2.shape[0]//2, i + kernel2.shape[0]//2):
                if j < 0 or j >= img_dim: continue
                A_small2[i, j] = kernel2[j - i + kernel2.shape[0]//2]
        self.U_small1, self.singulars_small1, self.V_small1 = torch.linalg.svd(A_small1)
        self.U_small2, self.singulars_small2, self.V_small2 = torch.linalg.svd(A_small2)
        ZERO = 3e-2
        self.singulars_small1[self.singulars_small1 < ZERO] = 0
        self.singulars_small2[self.singulars_small2 < ZERO] = 0
        #calculate the singular values of the big matrix
        self._singulars = torch.matmul(self.singulars_small1.reshape(img_dim, 1), self.singulars_small2.reshape(1, img_dim)).reshape(img_dim**2)
        #sort the big matrix singulars and save the permutation
        self._singulars, self._perm = self._singulars.sort(descending=True)
//  [python]
def forward(ctx, input):
        # Note: input is matrix of size m x n with m >= n.
        # Note: if above assumption is violated, the gradients
        # will be wrong.

        U, S, V = torch.svd(input, some=True)
        ctx.save_for_backward(U, S, V)
        return U, S, V
//  [python]
def forward(ctx, input):
        U, S, V = torch.linalg.svd(input, full_matrices=False)
        ctx.save_for_backward(U, S, V)
        return U, S, V
//  [python]
def __init__(self, dim, bias=True, identity_init=True):
        super().__init__()
        self.dim = dim

        _l = nn.Linear(dim, dim, bias)
        if identity_init:
            self.weight = nn.Parameter(torch.eye(dim))#[torch.randperm(dim)])
        else:
            UDV = torch.svd(_l.weight.data.t())
            self.weight = nn.Parameter(UDV[0])
            del UDV
        if bias:
            self.bias = nn.Parameter(_l.bias.data)
        else:
            self.bias = None
        del _l
//  [python]
def __init__(self, dim, bias=True, identity_init=True):
        super().__init__()
        self.dim = dim

        _l = nn.Linear(dim, dim, bias)
        if identity_init:
            self.weight = nn.Parameter(torch.eye(dim))#[torch.randperm(dim)])
        else:
            UDV = torch.linalg.svd(_l.weight.data.t())
            self.weight = nn.Parameter(UDV[0])
            del UDV
        if bias:
            self.bias = nn.Parameter(_l.bias.data)
        else:
            self.bias = None
        del _l