jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def project_matrix_to_unitary(W):
    if isinstance(W, np.ndarray):
        U, _, V = np.linalg.svd(W, full_matrices=True)
        U_refine = np.matmul(U, V)
    elif isinstance(W, torch.Tensor):
        U, S, V = torch.linalg.svd(W)
        U_refine = torch.matmul(U, V.transpose(-2, -1))
    else:
        raise NotImplementedError
    return U_refine
//  [python]
def qubitunitarystrict_matrix(params):
    """Compute unitary matrix for Qubitunitary strict gate.
        Strictly be the unitary.

    Args:
        params (torch.Tensor): The unitary matrix.

    Returns:
        torch.Tensor: The computed unitary matrix.

    """
    params.squeeze(0)
    mat = params
    U, Sigma, V = torch.svd(mat)
    return U.matmul(V)
//  [python]
def qubitunitarystrict_matrix(params):
    """Compute unitary matrix for Qubitunitary strict gate.
        Strictly be the unitary.

    Args:
        params (torch.Tensor): The unitary matrix.

    Returns:
        torch.Tensor: The computed unitary matrix.

    """
    params.squeeze(0)
    mat = params
    U, S, V = torch.linalg.svd(mat)
    return U.matmul(V)
//  [python]
def compute_similarity_transform(S1: torch.Tensor, S2: torch.Tensor) -> torch.Tensor:
    """
    Computes a similarity transform (sR, t) in a batched way that takes
    a set of 3D points S1 (B, N, 3) closest to a set of 3D points S2 (B, N, 3),
    where R is a 3x3 rotation matrix, t 3x1 translation, s scale.
    i.e. solves the orthogonal Procrutes problem.
    Args:
        S1 (torch.Tensor): First set of points of shape (B, N, 3).
        S2 (torch.Tensor): Second set of points of shape (B, N, 3).
    """

    batch_size = S1.shape[0]
    S1 = S1.permute(0, 2, 1)
    S2 = S2.permute(0, 2, 1)
    # 1. Remove mean.
    mu1 = S1.mean(dim=2, keepdim=True)
    mu2 = S2.mean(dim=2, keepdim=True)
    X1 = S1 - mu1
    X2 = S2 - mu2

    # 2. Compute variance of X1 used for scale.
    var1 = (X1**2).sum(dim=(1,2))

    # 3. The outer product of X1 and X2.
    K = torch.matmul(X1, X2.permute(0, 2, 1))

    # 4. Solution that Maximizes trace(R'K) is R=U*V', where U, V are singular vectors of K.
    U, s, V = torch.svd(K)
    Vh = V.permute(0, 2, 1)

    # Construct Z that fixes the orientation of R to get det(R)=1.
    Z = torch.eye(U.shape[1]).unsqueeze(0).repeat(batch_size, 1, 1)
    Z[:, -1, -1] *= torch.sign(torch.linalg.det(torch.matmul(U, Vh)))

    # Construct R.
    R = torch.matmul(torch.matmul(V, Z), U.permute(0, 2, 1))

    # 5. Recover scale.
    trace = torch.matmul(R, K).diagonal(offset=0, dim1=-1, dim2=-2).sum(dim=-1)
    scale = (trace / var1).unsqueeze(dim=-1).unsqueeze(dim=-1)

    # 6. Recover translation.
    t = mu2 - scale*torch.matmul(R, mu1)

    # 7. Error:
    S1_hat = scale*torch.matmul(R, S1) + t

    return R, t, S1_hat.permute(0, 2, 1)
//  [python]
def compute_similarity_transform(S1: torch.Tensor, S2: torch.Tensor) -> torch.Tensor:
    """
    Computes a similarity transform (sR, t) in a batched way that takes
    a set of 3D points S1 (B, N, 3) closest to a set of 3D points S2 (B, N, 3),
    where R is a 3x3 rotation matrix, t 3x1 translation, s scale.
    i.e. solves the orthogonal Procrutes problem.
    Args:
        S1 (torch.Tensor): First set of points of shape (B, N, 3).
        S2 (torch.Tensor): Second set of points of shape (B, N, 3).
    """

    batch_size = S1.shape[0]
    S1 = S1.permute(0, 2, 1)
    S2 = S2.permute(0, 2, 1)
    # 1. Remove mean.
    mu1 = S1.mean(dim=2, keepdim=True)
    mu2 = S2.mean(dim=2, keepdim=True)
    X1 = S1 - mu1
    X2 = S2 - mu2

    # 2. Compute variance of X1 used for scale.
    var1 = (X1**2).sum(dim=(1,2))

    # 3. The outer product of X1 and X2.
    K = torch.matmul(X1, X2.permute(0, 2, 1))

    # 4. Solution that Maximizes trace(R'K) is R=U*V', where U, V are singular vectors of K.
    U, s, V = torch.linalg.svd(K)
    Vh = V.permute(0, 2, 1)

    # Construct Z that fixes the orientation of R to get det(R)=1.
    Z = torch.eye(U.shape[1]).unsqueeze(0).repeat(batch_size, 1, 1)
    Z[:, -1, -1] *= torch.sign(torch.linalg.det(torch.matmul(U, Vh)))

    # Construct R.
    R = torch.matmul(torch.matmul(V, Z), U.permute(0, 2, 1))

    # 5. Recover scale.
    trace = torch.matmul(R, K).diagonal(offset=0, dim1=-1, dim2=-2).sum(dim=-1)
    scale = (trace / var1).unsqueeze(dim=-1).unsqueeze(dim=-1)

    # 6. Recover translation.
    t = mu2 - scale*torch.matmul(R, mu1)

    # 7. Error:
    S1_hat = scale*torch.matmul(R, S1) + t

    return R, t, S1_hat.permute(0, 2, 1)
//  [python]
def prox(self, x, step_size=None):
        """
        Projection operator on the Nuclear Norm constraint set.
        """
        U, S, V = torch.svd(x)
        # Project S on the alpha-L1 ball
        ball = L1Ball(self.alpha)

        S_proj = ball.prox(S.view(-1, S.size(-1))).view_as(S)

        VT = V.transpose(-2, -1)
        return torch.matmul(U, torch.matmul(torch.diag_embed(S_proj), VT))
//  [python]
def prox(self, x, step_size=None):
        """
        Projection operator on the Nuclear Norm constraint set.
        """
        U, S, V = torch.linalg.svd(x)
        # Project S on the alpha-L1 ball
        ball = L1Ball(self.alpha)

        S_proj = ball.prox(S.view(-1, S.size(-1))).view_as(S)

        VT = V.transpose(-2, -1)
        return torch.matmul(U, torch.matmul(torch.diag_embed(S_proj), VT))
//  [python]
def test_SVDGESDD_COMPLEX_random():
    
    def test_f_1(M):
        U,S,V= SVDGESDD_COMPLEX.apply(M)
        return torch.sum(S[0:1])

    def test_f_2(M):
        U,S,V= SVDGESDD_COMPLEX.apply(M)
        T= U @ V.conj().transpose(-2,-1)
        return T.norm()

    M= 25
    A= torch.rand((M, M), dtype=torch.complex128)
    U,S,V= torch.svd(A) 
    print(S)

    for split_scale in [10.0, 1.0, 0.1, 0.01, 0.]: 
        tot_scale=1000
        d0= torch.rand(M//2, dtype=torch.float64)
        splits= torch.rand(M//2, dtype=torch.float64)
        for i in range(M//2):
            S[2*i]= tot_scale*d0[i]
            S[2*i+1]= tot_scale*d0[i]+split_scale*splits[i]
        A= U * torch.diag(S) @ V.conj().transpose(-2,-1)
        A.requires_grad_()
        print(f"split_scale {split_scale}")
        print(S)

        assert(torch.autograd.gradcheck(test_f_1, A, eps=1e-6, atol=1e-4))
        assert(torch.autograd.gradcheck(test_f_2, A, eps=1e-6, atol=1e-4))
//  [python]
def test_SVDGESDD_COMPLEX_random():
    
    def test_f_1(M):
        U,S,V= SVDGESDD_COMPLEX.apply(M)
        return torch.sum(S[0:1])

    def test_f_2(M):
        U,S,V= SVDGESDD_COMPLEX.apply(M)
        T= U @ V.conj().transpose(-2,-1)
        return T.norm()

    M= 25
    A= torch.rand((M, M), dtype=torch.complex128)
    U,S,V= torch.linalg.svd(A) 
    print(S)

    for split_scale in [10.0, 1.0, 0.1, 0.01, 0.]: 
        tot_scale=1000
        d0= torch.rand(M//2, dtype=torch.float64)
        splits= torch.rand(M//2, dtype=torch.float64)
        for i in range(M//2):
            S[2*i]= tot_scale*d0[i]
            S[2*i+1]= tot_scale*d0[i]+split_scale*splits[i]
        A= U * torch.diag(S) @ V.conj().transpose(-2,-1)
        A.requires_grad_()
        print(f"split_scale {split_scale}")
        print(S)

        assert(torch.autograd.gradcheck(test_f_1, A, eps=1e-6, atol=1e-4))
        assert(torch.autograd.gradcheck(test_f_2, A, eps=1e-6, atol=1e-4))
//  [python]
def _evaluate_whls_zvp(self, weights, vlines):
        vlines = F.normalize(vlines, p=2, dim=-1)
        u, s, v = torch.svd(weights * vlines)
        return v[:, :, :, -1]
//  [python]
def _evaluate_whls_zvp(self, weights, vlines):
        vlines = F.normalize(vlines, p=2, dim=-1)
        u, s, v = torch.linalg.svd(weights * vlines)
        return v[:, :, :, -1]
//  [python]
def __init__(self, h):
        self.u, self.d, self.v = torch.svd(h, some=False)
        self.ut = self.u.t()
        self.vt = self.v.t()
        self.d[self.d < 1e-3] = 0
//  [python]
def __init__(self, h):
        self.u, self.s, self.v = torch.linalg.svd(h)
        self.ut = self.u.t()
        self.vt = self.v.t()
        self.s[self.s < 1e-3] = 0
//  [python]
def _get_zca_matrix(X, reg_coef=0.1):
    X_flat = X.reshape(X.shape[0], -1)
    cov = (X_flat.T @ X_flat) / X_flat.shape[0]
    reg_amount = reg_coef * torch.trace(cov) / cov.shape[0]
    u, s, _ = torch.svd(cov + reg_amount * torch.eye(cov.shape[0]).to(X.device))
    inv_sqrt_zca_eigs = s ** (-0.5)
    whitening_transform = torch.einsum("ij,j,kj->ik", u, inv_sqrt_zca_eigs, u)

    return whitening_transform
//  [python]
def _get_zca_matrix(X, reg_coef=0.1):
    X_flat = X.reshape(X.shape[0], -1)
    cov = (X_flat.T @ X_flat) / X_flat.shape[0]
    reg_amount = reg_coef * torch.trace(cov) / cov.shape[0]
    u, s, _ = torch.linalg.svd(cov + reg_amount * torch.eye(cov.shape[0]).to(X.device))
    inv_sqrt_zca_eigs = s ** (-0.5)
    whitening_transform = torch.einsum("ij,j,kj->ik", u, inv_sqrt_zca_eigs, u)

    return whitening_transform
//  [python]
def __init__(self, c, h, w):
        self.color_dim = c
        self.num_pixels = h * w
        self.u, self.d, self.v = torch.svd(torch.ones(1, c) / c, some=False)
        self.vt = self.v.t()
//  [python]
def __init__(self, c, h, w):
        self.color_dim = c
        self.num_pixels = h * w
        self.u, self.s, self.v = torch.linalg.svd(torch.ones(1, c) / c)
        self.vt = self.v.t()
//  [python]
def forward(self, input, alpha=1.):
        # return nn.functional.linear(input, torch.clamp(self.log_weight.exp(), min=0.0, max=1e10), bias=self.bias)

        # take the upper triangular part, and then set diagonal=1
        # return nn.functional.linear(input, torch.triu(self.weight,diagonal=1) + alpha * torch.eye(n=self.weight.shape[0], m=self.weight.shape[1]).cuda())

        # return nn.functional.linear(input, torch.eye(n=self.weight.shape[0], m=self.weight.shape[1])%de)

        # self.weight[torch.isnan(self.weight)] = 1e-5
        
        # self.weight[self.weight!=self.weight] = 1e-5 # replace nan with 1e-5
        U, S, V = torch.svd(self.weight+1e-5*torch.eye(n=self.weight.shape[0], m=self.weight.shape[1]).cuda())
        # S_fr = torch.clamp(torch.abs(S), min=1e-5, max=10.) * torch.sign(S)
        S_fr = torch.clamp(S, min=1e-5, max=10.)
        # print(S, S_fr)

        # print(S, S_fr)
        # weight = torch.matmul(torch.matmul(self.U, torch.eye(n=self.diag.shape[0]).cuda()), self.V.T)
        # print(torch.dist(self.weight, U @ torch.diag_embed(S) @ V.T))
        return nn.functional.linear(input, U @ torch.diag_embed(S_fr) @ V.T, bias=self.bias)
//  [python]
def forward(self, input, alpha=1.):
        # return nn.functional.linear(input, torch.clamp(self.log_weight.exp(), min=0.0, max=1e10), bias=self.bias)

        # take the upper triangular part, and then set diagonal=1
        # return nn.functional.linear(input, torch.triu(self.weight,diagonal=1) + alpha * torch.eye(n=self.weight.shape[0], m=self.weight.shape[1]).cuda())

        # return nn.functional.linear(input, torch.eye(n=self.weight.shape[0], m=self.weight.shape[1])%de)

        # self.weight[torch.isnan(self.weight)] = 1e-5
        
        # self.weight[self.weight!=self.weight] = 1e-5 # replace nan with 1e-5
        U, S, V = torch.linalg.svd(self.weight+1e-5*torch.eye(n=self.weight.shape[0], m=self.weight.shape[1]).cuda())
        # S_fr = torch.clamp(torch.abs(S), min=1e-5, max=10.) * torch.sign(S)
        S_fr = torch.clamp(S, min=1e-5, max=10.)
        # print(S, S_fr)

        # print(S, S_fr)
        # weight = torch.matmul(torch.matmul(self.U, torch.eye(n=self.diag.shape[0]).cuda()), self.V.T)
        # print(torch.dist(self.weight, U @ torch.diag_embed(S) @ V.T))
        return nn.functional.linear(input, U @ torch.diag_embed(S_fr) @ V.T, bias=self.bias)
//  [python]
def inv_sqrtm(A: torch.Tensor, eps: float = 1e-9):
    """Compute the inverse square-root of a positive definite matrix."""
    # Perform eigendecomposition of covariance matrix
    U, S, V = torch.svd(A)
    # Enforce positive definite by taking a torch max() with eps
    S = torch.clamp(S, min=eps)
    # S = torch.max(S, torch.tensor(eps, device=S.device))
    # Calculate inverse square-root
    inv_sqrt_S = torch.diag_embed(torch.pow(S, -0.5))
    # Calculate inverse square-root matrix
    B = torch.matmul(torch.matmul(U, inv_sqrt_S), V.transpose(-1, -2))
    return B
//  [python]
def inv_sqrtm(A: torch.Tensor, eps: float = 1e-9):
    """Compute the inverse square-root of a positive definite matrix."""
    # Perform eigendecomposition of covariance matrix
    U, S, V = torch.linalg.svd(A)
    # Enforce positive definite by taking a torch max() with eps
    S = torch.clamp(S, min=eps)
    # S = torch.max(S, torch.tensor(eps, device=S.device))
    # Calculate inverse square-root
    inv_sqrt_S = torch.diag_embed(torch.pow(S, -0.5))
    # Calculate inverse square-root matrix
    B = torch.matmul(torch.matmul(U, inv_sqrt_S), V.transpose(-1, -2))
    return B
//  [python]
def solve_rigid(self, x, y, w):
        """

        :param x: BxNxD
        :param y: BxNxD
        :param w: BxNx1
        :return:
        """
        B, N, D = x.shape[0], x.shape[1], x.shape[2]
        device = x.device
        sum_w = w.sum(1, keepdim=True)
        mu_x = (x * w).sum(1, keepdim=True) / sum_w
        mu_y = (y * w).sum(1, keepdim=True) / sum_w
        x_hat = x - mu_x
        wx_hat = x_hat * w
        y_hat = y - mu_y
        wy_hat = y_hat * w
        a = wy_hat.transpose(2, 1) @ wx_hat  # BxDxN @ BxNxD  BxDxD
        u, s, v = torch.svd(a)
        c = torch.ones(B, D).to(device)
        c[:, -1] = torch.det(u @ v)  #
        r = (u * (c[..., None])) @ v.transpose(2, 1)
        tr_atr = torch.diagonal(a.transpose(2, 1) @ r, dim1=-2, dim2=-1).sum(-1)
        tr_xtwx = torch.diagonal(wx_hat.transpose(2, 1) @ wx_hat, dim1=-2, dim2=-1).sum(
            -1
        )
        s = (
            (tr_atr / tr_xtwx)[..., None][..., None]
            if self.eval_scale_for_rigid
            else 1.0
        )
        t = mu_y - s * (r @ mu_x.transpose(2, 1)).transpose(2, 1)
        A = torch.cat([r.transpose(2, 1) * s, t], 1)
        X = torch.cat((x, torch.ones_like(x[:, :, :1])), dim=2)  # (B,N, D+1)
        return A, X @ A
//  [python]
def solve_rigid(self, x, y, w):
        """

        :param x: BxNxD
        :param y: BxNxD
        :param w: BxNx1
        :return:
        """
        B, N, D = x.shape[0], x.shape[1], x.shape[2]
        device = x.device
        sum_w = w.sum(1, keepdim=True)
        mu_x = (x * w).sum(1, keepdim=True) / sum_w
        mu_y = (y * w).sum(1, keepdim=True) / sum_w
        x_hat = x - mu_x
        wx_hat = x_hat * w
        y_hat = y - mu_y
        wy_hat = y_hat * w
        a = wy_hat.transpose(2, 1) @ wx_hat  # BxDxN @ BxNxD  BxDxD
        u, s, v = torch.linalg.svd(a)
        c = torch.ones(B, D).to(device)
        c[:, -1] = torch.det(u @ v)  #
        r = (u * (c[..., None])) @ v.transpose(2, 1)
        tr_atr = torch.diagonal(a.transpose(2, 1) @ r, dim1=-2, dim2=-1).sum(-1)
        tr_xtwx = torch.diagonal(wx_hat.transpose(2, 1) @ wx_hat, dim1=-2, dim2=-1).sum(
            -1
        )
        s = (
            (tr_atr / tr_xtwx)[..., None][..., None]
            if self.eval_scale_for_rigid
            else 1.0
        )
        t = mu_y - s * (r @ mu_x.transpose(2, 1)).transpose(2, 1)
        A = torch.cat([r.transpose(2, 1) * s, t], 1)
        X = torch.cat((x, torch.ones_like(x[:, :, :1])), dim=2)  # (B,N, D+1)
        return A, X @ A
//  [python]
def _svd_lowrank(A, q=6, niter=2, M=None):
    # type: (Tensor, Optional[int], Optional[int], Optional[Tensor]) -> Tuple[Tensor, Tensor, Tensor]
    q = 6 if q is None else q
    m, n = A.shape[-2:]
    matmul = _utils.matmul
    if M is None:
        M_t = None
    else:
        M_t = _utils.transpose(M)
    A_t = _utils.transpose(A)

    # Algorithm 5.1 in Halko et al 2009, slightly modified to reduce
    # the number conjugate and transpose operations
    if m < n:
        # computing the SVD approximation of a transpose in order to
        # keep B shape minimal
        Q = get_approximate_basis(A_t, q, niter=niter, M=M_t)
        Q_c = _utils.conjugate(Q)
        if M is None:
            B_t = matmul(A, Q_c)
        else:
            B_t = matmul(A, Q_c) - matmul(M, Q_c)
        U, S, V = torch.svd(B_t)
        V = Q.matmul(V)
    else:
        Q = get_approximate_basis(A, q, niter=niter, M=M)
        Q_c = _utils.conjugate(Q)
        if M is None:
            B = matmul(A_t, Q_c)
        else:
            B = matmul(A_t, Q_c) - matmul(M_t, Q_c)
        U, S, V = torch.svd(_utils.transpose(B))
        U = Q.matmul(U)

    return U, S, V
//  [python]
def _svd_lowrank(A, q=6, niter=2, M=None):
    # type: (Tensor, Optional[int], Optional[int], Optional[Tensor]) -> Tuple[Tensor, Tensor, Tensor]
    q = 6 if q is None else q
    m, n = A.shape[-2:]
    matmul = _utils.matmul
    if M is None:
        M_t = None
    else:
        M_t = _utils.transpose(M)
    A_t = _utils.transpose(A)

    # Algorithm 5.1 in Halko et al 2009, slightly modified to reduce
    # the number conjugate and transpose operations
    if m < n:
        # computing the SVD approximation of a transpose in order to
        # keep B shape minimal
        Q = get_approximate_basis(A_t, q, niter=niter, M=M_t)
        Q_c = _utils.conjugate(Q)
        if M is None:
            B_t = matmul(A, Q_c)
        else:
            B_t = matmul(A, Q_c) - matmul(M, Q_c)
        U, S, V = torch.linalg.svd(B_t)
        V = Q.matmul(V)
    else:
        Q = get_approximate_basis(A, q, niter=niter, M=M)
        Q_c = _utils.conjugate(Q)
        if M is None:
            B = matmul(A_t, Q_c)
        else:
            B = matmul(A_t, Q_c) - matmul(M_t, Q_c)
        U, S, V = torch.linalg.svd(_utils.transpose(B))
        U = Q.matmul(U)

    return U, S, V
//  [python]
def kabsch(pred, true):
    def rmsd(V, W, eps=1e-6):
        L = V.shape[0]
        return torch.sqrt(torch.sum((V-W)*(V-W)) / L + eps)
    def centroid(X):
        return X.mean(dim=-2, keepdim=True)

    cP = centroid(pred)
    cT = centroid(true)
    pred = pred - cP
    true = true - cT
    C = torch.matmul(pred.permute(1,0), true)
    V, S, W = torch.svd(C)
    d = torch.ones([3,3], device=pred.device)
    d[:,-1] = torch.sign(torch.det(V)*torch.det(W))
    U = torch.matmul(d*V, W.permute(1,0)) # (IB, 3, 3)
    rpred = torch.matmul(pred, U) # (IB, L*3, 3)
    rms = rmsd(rpred, true)
    return rms, U, cP, cT
//  [python]
def kabsch(pred, true):
    def rmsd(V, W, eps=1e-6):
        L = V.shape[0]
        return torch.sqrt(torch.sum((V-W)*(V-W)) / L + eps)
    def centroid(X):
        return X.mean(dim=-2, keepdim=True)

    cP = centroid(pred)
    cT = centroid(true)
    pred = pred - cP
    true = true - cT
    C = torch.matmul(pred.permute(1,0), true)
    U, S, V = torch.linalg.svd(C)
    d = torch.ones([3,3], device=pred.device)
    d[:,-1] = torch.sign(torch.det(U)*torch.det(V))
    W = torch.matmul(d*U, V.permute(1,0)) # (IB, 3, 3)
    rpred = torch.matmul(pred, W) # (IB, L*3, 3)
    rms = rmsd(rpred, true)
    return rms, W, cP, cT
//  [python]
def pca(x:Tensor, k=2):
    "Compute PCA of `x` with `k` dimensions."
    x = x-torch.mean(x,0)
    U,S,V = torch.svd(x.t())
    return torch.mm(x,U[:,:k])
//  [python]
def pca(x:Tensor, k=2):
    "Compute PCA of `x` with `k` dimensions."
    x = x-torch.mean(x,0)
    U,S,V = torch.linalg.svd(x.t())
    return torch.mm(x,U[:,:k])
//  [python]
def plot_state_singular_values(self, sample_id, color, linewidth=1, length=-1, show_title=True, title="",
                                   xticks=None, yticks=None, ylim=None, xlim=None, log10=False):
        """
        Plot states singular values
        :param sample_id:
        :param color:
        :param linewidth:
        :param ylim:
        :param length:
        :param show_title:
        :param title:
        :param log10:
        :return:
        """
        # State sample
        state_sample = self._esn_cell_states[sample_id]

        # Correlation matrix R
        R = torch.mm(state_sample.t(), state_sample) / state_sample.size(0)

        # SVD on state
        _, S, _ = torch.svd(R)

        # Log10?
        if log10:
            S = torch.log10(S)
        # end if

        # Learning PC energy
        if length != -1:
            plt.plot(S[:length], color=color, linewidth=linewidth)
        else:
            plt.plot(S, color=color, linewidth=linewidth)
        # end if

        # Title
        if show_title:
            plt.title(title)
        # end if

        # X labels
        if xlim is not None: plt.xlim(xlim)
        if xticks is not None: plt.xticks(xticks)

        # Y limits
        if ylim is not None: plt.ylim(ylim)
        if yticks is not None: plt.yticks(yticks)
//  [python]
def plot_state_singular_values(self, sample_id, color, linewidth=1, length=-1, show_title=True, title="",
                                   xticks=None, yticks=None, ylim=None, xlim=None, log10=False):
        """
        Plot states singular values
        :param sample_id:
        :param color:
        :param linewidth:
        :param ylim:
        :param length:
        :param show_title:
        :param title:
        :param log10:
        :return:
        """
        # State sample
        state_sample = self._esn_cell_states[sample_id]

        # Correlation matrix R
        R = torch.mm(state_sample.t(), state_sample) / state_sample.size(0)

        # SVD on state
        _, S, _ = torch.linalg.svd(R)

        # Log10?
        if log10:
            S = torch.log10(S)
        # end if

        # Learning PC energy
        if length != -1:
            plt.plot(S[:length], color=color, linewidth=linewidth)
        else:
            plt.plot(S, color=color, linewidth=linewidth)
        # end if

        # Title
        if show_title:
            plt.title(title)
        # end if

        # X labels
        if xlim is not None: plt.xlim(xlim)
        if xticks is not None: plt.xticks(xticks)

        # Y limits
        if ylim is not None: plt.ylim(ylim)
        if yticks is not None: plt.yticks(yticks)
//  [python]
def forward(self, A, cutoff, diagnostics):
        U, S, V = torch.svd(A)
        cutoff= torch.as_tensor(cutoff, dtype=S.dtype, device=S.device)
        self.save_for_backward(U, S, V, cutoff)
        return U, S, V
//  [python]
def forward(self, A, cutoff, diagnostics):
        U, S, V = torch.linalg.svd(A)
        cutoff= torch.as_tensor(cutoff, dtype=S.dtype, device=S.device)
        self.save_for_backward(U, S, V, cutoff)
        return U, S, V
//  [python]
def _so_maxpool(self, x):
        while len(x.data.shape ) < 4:
            x = torch.unsqueeze(x, 0)
        # x = x.double()
        batchSize, tupleLen, nFeat, dimFeat = x.data.shape 
        x = torch.reshape(x, (-1, dimFeat))
        x = torch.unsqueeze(x, -1)
        x = x.matmul(x.transpose(1, 2))

        x = torch.reshape(x, (batchSize, tupleLen, nFeat, dimFeat, dimFeat))
        x = torch.max(x, 2).values
        x = torch.reshape(x, (-1, dimFeat, dimFeat))
        if self.do_pe:
            # u_, s_, vh_ = torch.linalg.svd(x)
            # dist = torch.dist(x, u_ @ torch.diag_embed(s_) @ vh_)
            # dist_same = torch.allclose(x, u_ @ torch.diag_embed(s_) @ vh_)
            # s_alpha = torch.pow(s_, 0.5)
            # x = u_ @ torch.diag_embed(s_alpha) @ vh_
            x = x.double()
            u_, s_, v_ = torch.svd(x) # For pytorch versions < 1.9
            # dist = torch.dist(x, u_ @ torch.diag_embed(s_) @  v_.transpose(-2, -1))
            # dist_same = torch.allclose(x, u_ @ torch.diag_embed(s_) @  v_.transpose(-2, -1))
            s_alpha = torch.pow(s_, 0.5)
            x = u_ @ torch.diag_embed(s_alpha) @ v_.transpose(-2, -1)

        x = torch.reshape(x, (batchSize, tupleLen, dimFeat, dimFeat))
        return x
//  [python]
def _so_maxpool(self, x):
        while len(x.data.shape ) < 4:
            x = torch.unsqueeze(x, 0)
        # x = x.double()
        batchSize, tupleLen, nFeat, dimFeat = x.data.shape 
        x = torch.reshape(x, (-1, dimFeat))
        x = torch.unsqueeze(x, -1)
        x = x.matmul(x.transpose(1, 2))

        x = torch.reshape(x, (batchSize, tupleLen, nFeat, dimFeat, dimFeat))
        x = torch.max(x, 2).values
        x = torch.reshape(x, (-1, dimFeat, dimFeat))
        if self.do_pe:
            x = x.double()
            u_, s_, v_ = torch.linalg.svd(x)
            s_alpha = torch.pow(s_, 0.5)
            x = u_ @ torch.diag_embed(s_alpha) @ v_.transpose(-2, -1)

        x = torch.reshape(x, (batchSize, tupleLen, dimFeat, dimFeat))
        return x
//  [python]
def test_svd(self):
        linear_operator = self.create_linear_operator().requires_grad_(True)
        linear_operator_copy = linear_operator.clone().detach_().requires_grad_(True)
        evaluated = self.evaluate_linear_operator(linear_operator_copy)

        # Perform forward pass
        U_unsorted, S_unsorted, V_unsorted = torch.svd(linear_operator)
        U_unsorted, V_unsorted = U_unsorted.to_dense(), V_unsorted.to_dense()

        # since LinearOperator.svd does not sort the singular values, we do this here for the check
        S, idxr = torch.sort(S_unsorted, dim=-1, descending=True)
        idxr = idxr.unsqueeze(-2).expand(U_unsorted.shape)
        U = torch.gather(U_unsorted, dim=-1, index=idxr)
        V = torch.gather(V_unsorted, dim=-1, index=idxr)

        # compute expected result from full tensor
        U_actual, S_actual, V_actual = torch.svd(evaluated.double())
        U_actual = U_actual.to(dtype=evaluated.dtype)
        S_actual = S_actual.to(dtype=evaluated.dtype)
        V_actual = V_actual.to(dtype=evaluated.dtype)

        # Check forward pass
        self.assertAllClose(S, S_actual, rtol=1e-4, atol=1e-3)
        lt_from_svd = U @ torch.diag_embed(S) @ V.transpose(-1, -2)
        self.assertAllClose(lt_from_svd, evaluated, rtol=1e-4, atol=1e-3)

        # if there are repeated singular values, we'll skip checking the singular vectors
        U_abs, U_actual_abs = U.abs(), U_actual.abs()
        V_abs, V_actual_abs = V.abs(), V_actual.abs()
        any_svals_repeated = False
        for idx in itertools.product(*[range(b) for b in S_actual.shape[:-1]]):
            Si = S_actual[idx]
            if torch.unique(Si.detach()).shape[-1] == Si.shape[-1]:  # detach to avoid pytorch/pytorch#41389
                self.assertAllClose(U_abs[idx], U_actual_abs[idx], rtol=1e-4, atol=1e-3)
                self.assertAllClose(V_abs[idx], V_actual_abs[idx], rtol=1e-4, atol=1e-3)
            else:
                any_svals_repeated = True

        # Perform backward pass
        svd_grad = torch.randn_like(S)
        ((S * svd_grad).sum()).backward()
        ((S_actual * svd_grad).sum()).backward()

        # Check grads if there were no repeated singular values
        if not any_svals_repeated:
            for arg, arg_copy in zip(linear_operator.representation(), linear_operator_copy.representation()):
                if arg_copy.requires_grad and arg_copy.is_leaf and arg_copy.grad is not None:
                    self.assertAllClose(arg.grad, arg_copy.grad, rtol=1e-4, atol=1e-3)
//  [python]
def test_svd(self):
        linear_operator = self.create_linear_operator().requires_grad_(True)
        linear_operator_copy = linear_operator.clone().detach_().requires_grad_(True)
        evaluated = self.evaluate_linear_operator(linear_operator_copy)

        # Perform forward pass
        U_unsorted, S_unsorted, V_unsorted = torch.linalg.svd(linear_operator)
        U_unsorted, V_unsorted = U_unsorted.to_dense(), V_unsorted.to_dense()

        # since LinearOperator.svd does not sort the singular values, we do this here for the check
        S, idxr = torch.sort(S_unsorted, dim=-1, descending=True)
        idxr = idxr.unsqueeze(-2).expand(U_unsorted.shape)
        U = torch.gather(U_unsorted, dim=-1, index=idxr)
        V = torch.gather(V_unsorted, dim=-1, index=idxr)

        # compute expected result from full tensor
        U_actual, S_actual, V_actual = torch.linalg.svd(evaluated.double())
        U_actual = U_actual.to(dtype=evaluated.dtype)
        S_actual = S_actual.to(dtype=evaluated.dtype)
        V_actual = V_actual.to(dtype=evaluated.dtype)

        # Check forward pass
        self.assertAllClose(S, S_actual, rtol=1e-4, atol=1e-3)
        lt_from_svd = U @ torch.diag_embed(S) @ V.transpose(-1, -2)
        self.assertAllClose(lt_from_svd, evaluated, rtol=1e-4, atol=1e-3)

        # if there are repeated singular values, we'll skip checking the singular vectors
        U_abs, U_actual_abs = U.abs(), U_actual.abs()
        V_abs, V_actual_abs = V.abs(), V_actual.abs()
        any_svals_repeated = False
        for idx in itertools.product(*[range(b) for b in S_actual.shape[:-1]]):
            Si = S_actual[idx]
            if torch.unique(Si.detach()).shape[-1] == Si.shape[-1]:  # detach to avoid pytorch/pytorch#41389
                self.assertAllClose(U_abs[idx], U_actual_abs[idx], rtol=1e-4, atol=1e-3)
                self.assertAllClose(V_abs[idx], V_actual_abs[idx], rtol=1e-4, atol=1e-3)
            else:
                any_svals_repeated = True

        # Perform backward pass
        svd_grad = torch.randn_like(S)
        ((S * svd_grad).sum()).backward()
        ((S_actual * svd_grad).sum()).backward()

        # Check grads if there were no repeated singular values
        if not any_svals_repeated:
            for arg, arg_copy in zip(linear_operator.representation(), linear_operator_copy.representation()):
                if arg_copy.requires_grad and arg_copy.is_leaf and arg_copy.grad is not None:
                    self.assertAllClose(arg.grad, arg_copy.grad, rtol=1e-4, atol=1e-3)
//  [python]
def project_kd(p, x, keep_ambient=True):
    """ Project n points in dimension d onto 'direction' spanned by k ideal points
    p: (..., k, d) ideal points
    x: (..., n, d) points to project

    Returns:
    projection_1: (..., n, s) where s = d if keep_ambient==True otherwise s = k
    projection_2: same as projection_1. this is guaranteed to be the ideal point in the case k = 1
    p: the ideal points
    """

    if len(p.shape) < 2:
        p = p.unsqueeze(0)
    if len(x.shape) < 2:
        x = x.unsqueeze(0)
    k = p.size(-2)
    d = x.size(-1)
    assert d == p.size(-1)
    busemann_distances = busemann(x.unsqueeze(-2), p.unsqueeze(-3), keepdim=False)  # (..., n, k)
    c, r = busemann_to_horocycle(p.unsqueeze(-3), busemann_distances)  # (..., n, k, d) (..., n, k)
    c, r, ortho = sphere_intersections(c, r)  # (..., n, d) (..., n) (..., n, d, k-1)
    # we are looking for a vector spanned by the k ideal points, orthogonal to k-1 given vectors
    # i.e. x @ p @ ortho = 0
    if ortho is None:
        direction = torch.ones_like(busemann_distances)  # (..., n, k)
    else:
        a = torch.matmul(p.unsqueeze(-3), ortho)  # (..., n, k, k-1) = (..., n, k, d) @ (..., n, d, k-1)
        u, s, v = torch.svd(a, some=False)  # a = u s v^T
        direction = u[..., -1]  # (..., n, k)
    direction = direction @ p  # (..., n, d)
    direction = direction / torch.norm(direction, dim=-1, keepdim=True).clamp_min(MIN_NORM)

    projection_1 = c - r.unsqueeze(-1) * direction
    projection_2 = c + r.unsqueeze(-1) * direction
    if not keep_ambient:
        _, _, v = torch.svd(p, some=False)  # P = USV^T => PV = US so last d-k columns of PV are 0
        projection_1 = (projection_1 @ v)[..., :k]
        projection_2 = (projection_2 @ v)[..., :k]
        p = (p @ v)[..., :k]

    return projection_1, projection_2, p
//  [python]
def project_kd(p, x, keep_ambient=True):
    """ Project n points in dimension d onto 'direction' spanned by k ideal points
    p: (..., k, d) ideal points
    x: (..., n, d) points to project

    Returns:
    projection_1: (..., n, s) where s = d if keep_ambient==True otherwise s = k
    projection_2: same as projection_1. this is guaranteed to be the ideal point in the case k = 1
    p: the ideal points
    """

    if len(p.shape) < 2:
        p = p.unsqueeze(0)
    if len(x.shape) < 2:
        x = x.unsqueeze(0)
    k = p.size(-2)
    d = x.size(-1)
    assert d == p.size(-1)
    busemann_distances = busemann(x.unsqueeze(-2), p.unsqueeze(-3), keepdim=False)  # (..., n, k)
    c, r = busemann_to_horocycle(p.unsqueeze(-3), busemann_distances)  # (..., n, k, d) (..., n, k)
    c, r, ortho = sphere_intersections(c, r)  # (..., n, d) (..., n) (..., n, d, k-1)
    # we are looking for a vector spanned by the k ideal points, orthogonal to k-1 given vectors
    # i.e. x @ p @ ortho = 0
    if ortho is None:
        direction = torch.ones_like(busemann_distances)  # (..., n, k)
    else:
        a = torch.matmul(p.unsqueeze(-3), ortho)  # (..., n, k, k-1) = (..., n, k, d) @ (..., n, d, k-1)
        u, s, v = torch.linalg.svd(a)  # a = u s v^T
        direction = u[..., -1]  # (..., n, k)
    direction = direction @ p  # (..., n, d)
    direction = direction / torch.norm(direction, dim=-1, keepdim=True).clamp_min(MIN_NORM)

    projection_1 = c - r.unsqueeze(-1) * direction
    projection_2 = c + r.unsqueeze(-1) * direction
    if not keep_ambient:
        U, S, V = torch.linalg.svd(p)  # P = USV^T => PV = US so last d-k columns of PV are 0
        projection_1 = (projection_1 @ V)[..., :k]
        projection_2 = (projection_2 @ V)[..., :k]
        p = (p @ V)[..., :k]

    return projection_1, projection_2, p
//  [python]
def Coles(ndatas, K, lam, lowrank=5):
    ndatas = QRreduction(ndatas)#to speed up svd,not necessary for the performance
    u, _, _ = torch.svd(ndatas)
    u = u[:, :, :lowrank]
    W = torch.abs(u.matmul(torch.transpose(u,dim0=2,dim1=1)))
    for i in range(W.shape[0]):
        W[i,:,:].squeeze().fill_diagonal_(0)
    isqrt_diag = 1. / torch.sqrt(1e-4 + torch.sum(W, dim=-1,keepdim=True))
    W = W * isqrt_diag * torch.transpose(isqrt_diag,dim0=2,dim1=1)
    W1 = torch.ones_like(W)
    for i in range(W1.shape[0]):
        W1[i,:,:].squeeze().fill_diagonal_(0)
    isqrt_diag = 1. / torch.sqrt(1e-4 + torch.sum(W1, dim=-1,keepdim=True))
    W1 = W1 * isqrt_diag * torch.transpose(isqrt_diag,dim0=2,dim1=1)
    lapReg1 = torch.transpose(ndatas,dim0=2,dim1=1).matmul(lam * W1 - W).matmul(ndatas)
    e, v1 = torch.linalg.eigh(lapReg1)
    return ndatas.matmul(v1[:, :, :K])
//  [python]
def Coles(ndatas, K, lam, lowrank=5):
    ndatas = QRreduction(ndatas)#to speed up svd,not necessary for the performance
    u, _, _ = torch.linalg.svd(ndatas)
    u = u[:, :, :lowrank]
    W = torch.abs(u.matmul(torch.transpose(u,dim0=2,dim1=1)))
    for i in range(W.shape[0]):
        W[i,:,:].squeeze().fill_diagonal_(0)
    isqrt_diag = 1. / torch.sqrt(1e-4 + torch.sum(W, dim=-1,keepdim=True))
    W = W * isqrt_diag * torch.transpose(isqrt_diag,dim0=2,dim1=1)
    W1 = torch.ones_like(W)
    for i in range(W1.shape[0]):
        W1[i,:,:].squeeze().fill_diagonal_(0)
    isqrt_diag = 1. / torch.sqrt(1e-4 + torch.sum(W1, dim=-1,keepdim=True))
    W1 = W1 * isqrt_diag * torch.transpose(isqrt_diag,dim0=2,dim1=1)
    lapReg1 = torch.transpose(ndatas,dim0=2,dim1=1).matmul(lam * W1 - W).matmul(ndatas)
    e, v1 = torch.linalg.eigh(lapReg1)
    return ndatas.matmul(v1[:, :, :K])
//  [python]
def mendoncaError(F,f):

    M = F.shape[0]

    # create intrinsic matrix
    A = torch.zeros((3,3))
    A[0,0] = f
    A[1,1] = f
    A[2,2] = 1

    # stack for batch matrix multiplication
    A_t = torch.stack(M*[A.T])
    A_ = torch.stack(M*[A])

    # get essential matrix
    EM = torch.bmm(torch.bmm(A_t,F),A_)

    # get singular values
    u,d,v = torch.svd(EM)
    r = d[:,0]
    s = d[:,1]

    # cost function
    cost = (1/M) * (r-s) / (r+s)

    return torch.sum(cost)
//  [python]
def mendoncaError(F,f):

    M = F.shape[0]

    # create intrinsic matrix
    A = torch.zeros((3,3))
    A[0,0] = f
    A[1,1] = f
    A[2,2] = 1

    # stack for batch matrix multiplication
    A_t = torch.stack(M*[A.T])
    A_ = torch.stack(M*[A])

    # get essential matrix
    EM = torch.bmm(torch.bmm(A_t,F),A_)

    # get singular values
    u,d,v = torch.linalg.svd(EM)
    r = d[:,0]
    s = d[:,1]

    # cost function
    cost = (1/M) * (r-s) / (r+s)

    return torch.sum(cost)
//  [python]
def __init__(self, dims_in, dims_c=None,
                 correction_interval: int = 256,
                 clamp: float = 5.):
        '''
        Args:

          correction_interval: After this many gradient steps, the matrix is
            projected back to the Stiefel manifold to make it perfectly orthogonal.
          clamp: clamps the log scaling for stability. Corresponds to
            :math:`alpha` above.
        '''
        super().__init__(dims_in, dims_c)
        self.width = dims_in[0][0]
        self.clamp = clamp

        self.correction_interval = correction_interval
        self.back_counter = np.random.randint(0, correction_interval) // 2

        self.weights = torch.randn(self.width, self.width)
        self.weights = self.weights + self.weights.t()
        self.weights, S, V = torch.svd(self.weights)

        self.weights = nn.Parameter(self.weights)

        self.bias = nn.Parameter(0.05 * torch.randn(1, self.width))
        self.scaling = nn.Parameter(0.02 * torch.randn(1, self.width))

        self.register_backward_hook(correct_weights)
//  [python]
def __init__(self, dims_in, dims_c=None,
                 correction_interval: int = 256,
                 clamp: float = 5.):
        '''
        Args:

          correction_interval: After this many gradient steps, the matrix is
            projected back to the Stiefel manifold to make it perfectly orthogonal.
          clamp: clamps the log scaling for stability. Corresponds to
            :math:`alpha` above.
        '''
        super().__init__(dims_in, dims_c)
        self.width = dims_in[0][0]
        self.clamp = clamp

        self.correction_interval = correction_interval
        self.back_counter = np.random.randint(0, correction_interval) // 2

        self.weights = torch.randn(self.width, self.width)
        self.weights = self.weights + self.weights.t()
        self.weights, S, V = torch.linalg.svd(self.weights)

        self.weights = nn.Parameter(self.weights)

        self.bias = nn.Parameter(0.05 * torch.randn(1, self.width))
        self.scaling = nn.Parameter(0.02 * torch.randn(1, self.width))

        self.register_backward_hook(correct_weights)
//  [python]
def test_SVDGESDD_COMPLEX_random():
    
    def test_f_1(M):
        U,S,V= SVDGESDD.apply(M)
        return torch.sum(S[0:1])

    def test_f_2(M):
        U,S,V= SVDGESDD.apply(M)
        T= U @ V.conj().transpose(-2,-1)
        return T.norm()

    M= 25
    A= torch.rand((M, M), dtype=torch.complex128)
    U,S,V= torch.svd(A) 
    print(S)

    for split_scale in [10.0, 1.0, 0.1, 0.01, 0.]: 
        tot_scale=1000
        d0= torch.rand(M//2, dtype=torch.float64)
        splits= torch.rand(M//2, dtype=torch.float64)
        for i in range(M//2):
            S[2*i]= tot_scale*d0[i]
            S[2*i+1]= tot_scale*d0[i]+split_scale*splits[i]
        A= U * torch.diag(S) @ V.conj().transpose(-2,-1)
        A.requires_grad_()
        print(f"split_scale {split_scale}")
        print(S)

        assert(torch.autograd.gradcheck(test_f_1, A, eps=1e-6, atol=1e-4))
        assert(torch.autograd.gradcheck(test_f_2, A, eps=1e-6, atol=1e-4))
//  [python]
def test_SVDGESDD_COMPLEX_random():
    
    def test_f_1(M):
        U,S,V= SVDGESDD.apply(M)
        return torch.sum(S[0:1])

    def test_f_2(M):
        U,S,V= SVDGESDD.apply(M)
        T= U @ V.conj().transpose(-2,-1)
        return T.norm()

    M= 25
    A= torch.rand((M, M), dtype=torch.complex128)
    U,S,V= torch.linalg.svd(A) 
    print(S)

    for split_scale in [10.0, 1.0, 0.1, 0.01, 0.]: 
        tot_scale=1000
        d0= torch.rand(M//2, dtype=torch.float64)
        splits= torch.rand(M//2, dtype=torch.float64)
        for i in range(M//2):
            S[2*i]= tot_scale*d0[i]
            S[2*i+1]= tot_scale*d0[i]+split_scale*splits[i]
        A= U * torch.diag(S) @ V.conj().transpose(-2,-1)
        A.requires_grad_()
        print(f"split_scale {split_scale}")
        print(S)

        assert(torch.autograd.gradcheck(test_f_1, A, eps=1e-6, atol=1e-4))
        assert(torch.autograd.gradcheck(test_f_2, A, eps=1e-6, atol=1e-4))
//  [python]
def getExtrinsics_(x_c,p_w):
    c_center = torch.mean(x_c,dim=1)
    w_center = torch.mean(p_w,dim=1)

    # center the  3d shapes
    x_c_centered = x_c - c_center.unsqueeze(1)
    p_w_centered = p_w - w_center.unsqueeze(1)

    # create martrix to solve
    M = x_c.shape[0]
    Matrix = torch.bmm(x_c_centered.permute(0,2,1),p_w_centered)

    u,d,v = torch.svd(Matrix)

    #R = torch.bmm(u,v.permute(0,2,1))
    R = torch.bmm(u,v.permute(0,2,1))

    # solve T using centers
    #rot_w_center = torch.bmm(R,rep_w_center)

    # T = c_center - w_center.squeeze()
    T = c_center

    return R,T
//  [python]
def getExtrinsics_(x_c,p_w):
    c_center = torch.mean(x_c,dim=1)
    w_center = torch.mean(p_w,dim=1)

    # center the  3d shapes
    x_c_centered = x_c - c_center.unsqueeze(1)
    p_w_centered = p_w - w_center.unsqueeze(1)

    # create martrix to solve
    M = x_c.shape[0]
    Matrix = torch.bmm(x_c_centered.permute(0,2,1),p_w_centered)

    u,d,v = torch.linalg.svd(Matrix)

    #R = torch.bmm(u,v.permute(0,2,1))
    R = torch.bmm(u,v.permute(0,2,1))

    # solve T using centers
    #rot_w_center = torch.bmm(R,rep_w_center)

    # T = c_center - w_center.squeeze()
    T = c_center

    return R,T
//  [python]
def EPnP_(x_i,x_w,K):
    b = x_i.shape[0]

    # get control points
    c_w = getControlPoints(x_w).to(x_i.device).unsqueeze(0).repeat(b,1,1)

    # solve alphas
    alphas = solveAlphas_(x_w,c_w)

    # create matrix
    Matrix = setupM_(alphas,x_i,K)

    # get svd
    u,d,v = torch.svd(Matrix)
    km = v[:,:,-4:]

    # solve N = 1
    beta_n1 = torch.zeros((b,4)).to(x_i.device)
    beta_n1[:,3] = 1
    c_c_n1 = km[:,:,-1].reshape((b,4,3)).permute(0,2,1)
    s = scaleControlPoints_(c_c_n1,c_w,alphas,x_w)
    scaled_betas = beta_n1 * s.unsqueeze(1)

    # solve extrinsics
    Xc, R, T = optimizeGN_(km,c_w,scaled_betas,alphas,x_w,x_i)

    return Xc, R, T
//  [python]
def EPnP_(x_i,x_w,K):
    b = x_i.shape[0]

    # get control points
    c_w = getControlPoints(x_w).to(x_i.device).unsqueeze(0).repeat(b,1,1)

    # solve alphas
    alphas = solveAlphas_(x_w,c_w)

    # create matrix
    Matrix = setupM_(alphas,x_i,K)

    # get svd
    u,d,v = torch.linalg.svd(Matrix)
    km = v[:,:,-4:]

    # solve N = 1
    beta_n1 = torch.zeros((b,4)).to(x_i.device)
    beta_n1[:,3] = 1
    c_c_n1 = km[:,:,-1].reshape((b,4,3)).permute(0,2,1)
    s = scaleControlPoints_(c_c_n1,c_w,alphas,x_w)
    scaled_betas = beta_n1 * s.unsqueeze(1)

    # solve extrinsics
    Xc, R, T = optimizeGN_(km,c_w,scaled_betas,alphas,x_w,x_i)

    return Xc, R, T
//  [python]
def getExtrinsics(x_c,p_w):
    c_center = torch.mean(x_c,dim=1)
    w_center = torch.mean(p_w,dim=0)

    # center the  3d shapes
    x_c_centered = x_c - c_center.unsqueeze(1)
    p_w_centered = p_w - w_center.unsqueeze(0)

    # create martrix to solve
    M = x_c.shape[0]
    if p_w.is_cuda:
        Matrix = torch.zeros((M,3,3)).cuda()
    else:
        Matrix = torch.zeros((M,3,3))
    for i in range(M):
        m1 = x_c_centered[i].unsqueeze(2)
        m2 = p_w_centered.unsqueeze(1)
        tmpM = torch.bmm(m1,m2)
        Matrix[i] = torch.sum(tmpM,dim=0)

    u,d,v = torch.svd(Matrix)

    #R = torch.bmm(u,v.permute(0,2,1))
    R = torch.bmm(u,v.permute(0,2,1))

    # solve T using centers
    rep_w_center = torch.stack(M * [w_center]).unsqueeze(2)
    rot_w_center = torch.bmm(R,rep_w_center)
    T = c_center - rot_w_center.squeeze()

    return R,T
//  [python]
def getExtrinsics(x_c,p_w):
    c_center = torch.mean(x_c,dim=1)
    w_center = torch.mean(p_w,dim=0)

    # center the  3d shapes
    x_c_centered = x_c - c_center.unsqueeze(1)
    p_w_centered = p_w - w_center.unsqueeze(0)

    # create martrix to solve
    M = x_c.shape[0]
    if p_w.is_cuda:
        Matrix = torch.zeros((M,3,3)).cuda()
    else:
        Matrix = torch.zeros((M,3,3))
    for i in range(M):
        m1 = x_c_centered[i].unsqueeze(2)
        m2 = p_w_centered.unsqueeze(1)
        tmpM = torch.bmm(m1,m2)
        Matrix[i] = torch.sum(tmpM,dim=0)

    u,d,v = torch.linalg.svd(Matrix)

    #R = torch.bmm(u,v.permute(0,2,1))
    R = torch.bmm(u,v.permute(0,2,1))

    # solve T using centers
    rep_w_center = torch.stack(M * [w_center]).unsqueeze(2)
    rot_w_center = torch.bmm(R,rep_w_center)
    T = c_center - rot_w_center.squeeze()

    return R,T
//  [python]
def EPnP_single(x_img,x_w,K):
    M = x_img.shape[0]
    N = x_img.shape[2]
    f = K[:,0,0]
    px = K[:,0,2]
    py = K[:,1,2]

    # get control points
    c_w = getControlPoints(x_w)

    # solve alphas
    alphas = solveAlphas(x_w,c_w)
    if x_img.is_cuda:
        alphas = alphas.cuda()
    Matrix = setupM_single(alphas,x_img.permute(0,2,1),px,py,f)

    if ~torch.all(Matrix == Matrix):
        print(Matrix)
        print(f)
        print(x_w)
        return

    u,d,v = torch.svd(Matrix)
    km = v[:,:,-4:]

    # solve N=1
    beta_n1 = torch.zeros((M,4))
    if x_img.is_cuda:
        beta_n1 = beta_n1.cuda()
    beta_n1[:,3] = 1
    c_c_n1 = km[:,:,-1].reshape((M,4,3)).permute(0,2,1)
    _, x_c_n1,s1 = scaleControlPoints(c_c_n1,c_w[:3,:],alphas,x_w)
    mask1 = s1 == s1
    Rn1, Tn1 = getExtrinsics(x_c_n1[mask1],x_w)

    return km, c_w, beta_n1 / s1.unsqueeze(1), alphas
//  [python]
def EPnP_single(x_img,x_w,K):
    M = x_img.shape[0]
    N = x_img.shape[2]
    f = K[:,0,0]
    px = K[:,0,2]
    py = K[:,1,2]

    # get control points
    c_w = getControlPoints(x_w)

    # solve alphas
    alphas = solveAlphas(x_w,c_w)
    if x_img.is_cuda:
        alphas = alphas.cuda()
    Matrix = setupM_single(alphas,x_img.permute(0,2,1),px,py,f)

    if ~torch.all(Matrix == Matrix):
        print(Matrix)
        print(f)
        print(x_w)
        return

    u,s,v = torch.linalg.svd(Matrix)
    km = v[:,:,-4:]

    # solve N=1
    beta_n1 = torch.zeros((M,4))
    if x_img.is_cuda:
        beta_n1 = beta_n1.cuda()
    beta_n1[:,3] = 1
    c_c_n1 = km[:,:,-1].reshape((M,4,3)).permute(0,2,1)
    _, x_c_n1,s1 = scaleControlPoints(c_c_n1,c_w[:3,:],alphas,x_w)
    mask1 = s1 == s1
    Rn1, Tn1 = getExtrinsics(x_c_n1[mask1],x_w)

    return km, c_w, beta_n1 / s1.unsqueeze(1), alphas
//  [python]
def __init__(self, kernel, channels, img_dim, device):
        self.img_dim = img_dim
        self.channels = channels
        # build 1D conv matrix
        H_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
                H_small[i, j] = kernel[j - i + kernel.shape[0] // 2]
        # torch.cholesky(H_small) # raise error if H_small is not positive definite
        self.H_small = H_small # positive definite matrix
        # get the evd of the 1D conv
        self.U_small, self.singulars_small, _ = torch.svd(H_small, some=False)
        self.V_small = self.U_small # V should equal U since H is symmetric
        
        ZERO = 3e-2
        # truncate
        self.singulars_small[self.singulars_small < ZERO] = ZERO
        # calculate the singular values of the big matrix using Kronecker product
        self._singulars = torch.matmul(self.singulars_small.reshape(img_dim, 1),
                                       self.singulars_small.reshape(1, img_dim)).reshape(img_dim ** 2)
        self._singulars[self._singulars > 1] = 1
        print("contains zero?", torch.any(self._singulars == 0))
        # raise NotImplementedError
        # sort the big matrix singulars and save the permutation
        self._singulars, self._perm = self._singulars.sort(descending=True)