jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def extract_conv(
    weight: Union[torch.Tensor, nn.Parameter],
    mode="fixed",
    mode_param=0,
    device="cpu",
    is_cp=False,
) -> Tuple[nn.Parameter, nn.Parameter]:
    weight = weight.to(device)
    out_ch, in_ch, kernel_size, _ = weight.shape

    U, S, Vh = torch.linalg.svd(weight.reshape(out_ch, -1))

    if mode == "full":
        return weight, "full"
    elif mode == "fixed":
        lora_rank = mode_param
    elif mode == "threshold":
        assert mode_param >= 0
        lora_rank = torch.sum(S > mode_param)
    elif mode == "ratio":
        assert 1 >= mode_param >= 0
        min_s = torch.max(S) * mode_param
        lora_rank = torch.sum(S > min_s)
    elif mode == "quantile" or mode == "percentile":
        assert 1 >= mode_param >= 0
        s_cum = torch.cumsum(S, dim=0)
        min_cum_sum = mode_param * torch.sum(S)
        lora_rank = torch.sum(s_cum < min_cum_sum)
    else:
        raise NotImplementedError(
            'Extract mode should be "fixed", "threshold", "ratio" or "quantile"'
        )
    lora_rank = max(1, lora_rank)
    lora_rank = min(out_ch, in_ch, lora_rank)
    if lora_rank >= out_ch / 2 and not is_cp:
        return weight, "full"

    U = U[:, :lora_rank]
    S = S[:lora_rank]
    U = U @ torch.diag(S).to(device)
    Vh = Vh[:lora_rank, :]

    diff = (weight - (U @ Vh).reshape(out_ch, in_ch, kernel_size, kernel_size)).detach()
    extract_weight_A = Vh.reshape(lora_rank, in_ch, kernel_size, kernel_size).detach()
    extract_weight_B = U.reshape(out_ch, lora_rank, 1, 1).detach()
    del U, S, Vh, weight
    return (extract_weight_A, extract_weight_B, diff), "low rank"
//  [python]
def extract_conv(
    weight: Union[torch.Tensor, nn.Parameter],
    mode="fixed",
    mode_param=0,
    device="cpu",
    is_cp=False,
) -> Tuple[nn.Parameter, nn.Parameter]:
    weight = weight.to(device)
    out_ch, in_ch, kernel_size, _ = weight.shape

    U, S, Vh = linalg.svd(weight.reshape(out_ch, -1))

    if mode == "full":
        return weight, "full"
    elif mode == "fixed":
        lora_rank = mode_param
    elif mode == "threshold":
        assert mode_param >= 0
        lora_rank = torch.sum(S > mode_param)
    elif mode == "ratio":
        assert 1 >= mode_param >= 0
        min_s = torch.max(S) * mode_param
        lora_rank = torch.sum(S > min_s)
    elif mode == "quantile" or mode == "percentile":
        assert 1 >= mode_param >= 0
        s_cum = torch.cumsum(S, dim=0)
        min_cum_sum = mode_param * torch.sum(S)
        lora_rank = torch.sum(s_cum < min_cum_sum)
    else:
        raise NotImplementedError(
            'Extract mode should be "fixed", "threshold", "ratio" or "quantile"'
        )
    lora_rank = max(1, lora_rank)
    lora_rank = min(out_ch, in_ch, lora_rank)
    if lora_rank >= out_ch / 2 and not is_cp:
        return weight, "full"

    U = U[:, :lora_rank]
    S = S[:lora_rank]
    U = U @ torch.diag(S).to(device)
    Vh = Vh[:lora_rank, :]

    diff = (weight - (U @ Vh).reshape(out_ch, in_ch, kernel_size, kernel_size)).detach()
    extract_weight_A = Vh.reshape(lora_rank, in_ch, kernel_size, kernel_size).detach()
    extract_weight_B = U.reshape(out_ch, lora_rank, 1, 1).detach()
    del U, S, Vh, weight
    return (extract_weight_A, extract_weight_B, diff), "low rank"
//  [python]
def extract_linear(
    weight: Union[torch.Tensor, nn.Parameter],
    mode="fixed",
    mode_param=0,
    device="cpu",
) -> Tuple[nn.Parameter, nn.Parameter]:
    weight = weight.to(device)
    out_ch, in_ch = weight.shape

    U, S, Vh = torch.linalg.svd(weight)

    if mode == "full":
        return weight, "full"
    elif mode == "fixed":
        lora_rank = mode_param
    elif mode == "threshold":
        assert mode_param >= 0
        lora_rank = torch.sum(S > mode_param)
    elif mode == "ratio":
        assert 1 >= mode_param >= 0
        min_s = torch.max(S) * mode_param
        lora_rank = torch.sum(S > min_s)
    elif mode == "quantile" or mode == "percentile":
        assert 1 >= mode_param >= 0
        s_cum = torch.cumsum(S, dim=0)
        min_cum_sum = mode_param * torch.sum(S)
        lora_rank = torch.sum(s_cum < min_cum_sum)
    else:
        raise NotImplementedError(
            'Extract mode should be "fixed", "threshold", "ratio" or "quantile"'
        )
    lora_rank = max(1, lora_rank)
    lora_rank = min(out_ch, in_ch, lora_rank)
    if lora_rank >= out_ch / 2:
        return weight, "full"

    U = U[:, :lora_rank]
    S = S[:lora_rank]
    U = U @ torch.diag(S).to(device)
    Vh = Vh[:lora_rank, :]

    diff = (weight - U @ Vh).detach()
    extract_weight_A = Vh.reshape(lora_rank, in_ch).detach()
    extract_weight_B = U.reshape(out_ch, lora_rank).detach()
    del U, S, Vh, weight
    return (extract_weight_A, extract_weight_B, diff), "low rank"
//  [python]
def extract_linear(
    weight: Union[torch.Tensor, nn.Parameter],
    mode="fixed",
    mode_param=0,
    device="cpu",
) -> Tuple[nn.Parameter, nn.Parameter]:
    weight = weight.to(device)
    out_ch, in_ch = weight.shape

    U, S, Vh = linalg.svd(weight)

    if mode == "full":
        return weight, "full"
    elif mode == "fixed":
        lora_rank = mode_param
    elif mode == "threshold":
        assert mode_param >= 0
        lora_rank = torch.sum(S > mode_param)
    elif mode == "ratio":
        assert 1 >= mode_param >= 0
        min_s = torch.max(S) * mode_param
        lora_rank = torch.sum(S > min_s)
    elif mode == "quantile" or mode == "percentile":
        assert 1 >= mode_param >= 0
        s_cum = torch.cumsum(S, dim=0)
        min_cum_sum = mode_param * torch.sum(S)
        lora_rank = torch.sum(s_cum < min_cum_sum)
    else:
        raise NotImplementedError(
            'Extract mode should be "fixed", "threshold", "ratio" or "quantile"'
        )
    lora_rank = max(1, lora_rank)
    lora_rank = min(out_ch, in_ch, lora_rank)
    if lora_rank >= out_ch / 2:
        return weight, "full"

    U = U[:, :lora_rank]
    S = S[:lora_rank]
    U = U @ torch.diag(S).to(device)
    Vh = Vh[:lora_rank, :]

    diff = (weight - U @ Vh).detach()
    extract_weight_A = Vh.reshape(lora_rank, in_ch).detach()
    extract_weight_B = U.reshape(out_ch, lora_rank).detach()
    del U, S, Vh, weight
    return (extract_weight_A, extract_weight_B, diff), "low rank"
//  [python]
def fit(self, features, logits, targets):
        self.mean = features.mean(0, keepdim=True)
        U, S, V = torch.svd(features - self.mean)
        self.certificates = V[:, -self.k:]
//  [python]
def fit(self, features, logits, targets):
        self.mean = features.mean(0, keepdim=True)
        _, _, Vt = torch.linalg.svd(features - self.mean)
        self.certificates = Vt[:, -self.k:]
//  [python]
def fit(self, X):
        n, d = X.size()
        if self.n_components is not None:
            d = min(self.n_components, d)
        self.register_buffer("mean_", X.mean(0, keepdim=True))
        Z = X - self.mean_ # center
        U, S, Vh = torch.svd(Z)
        Vt = Vh.t()
        U, Vt = svd_flip(U, Vt)
        self.register_buffer("components_", Vt[:d])
        return self
//  [python]
def fit(self, X):
        n, d = X.size()
        if self.n_components is not None:
            d = min(self.n_components, d)
        self.register_buffer("mean_", X.mean(0, keepdim=True))
        Z = X - self.mean_ # center
        U, S, Vh = torch.linalg.svd(Z, full_matrices=False)
        Vt = Vh
        U, Vt = svd_flip(U, Vt)
        self.register_buffer("components_", Vt[:d])
        return self
//  [python]
def tt_svd(A, ranks):
    warnings.warn('Use tntorch instead for a better approximation')
    if not torch.is_tensor(A):
        raise ValueError('Operand is not a tensor')
    if type(ranks) in (list, tuple):
        tt_ranks_max = get_tt_ranks(list(A.shape))
        if A.dim() + 1 != len(ranks) or ranks[0] != ranks[-1] != 1 or \
                any(a > b for a, b in zip(ranks, tt_ranks_max)):
            raise ValueError('Invalid ranks')
    elif type(ranks) is int:
        ranks = get_tt_ranks(A.shape, max_rank=ranks)
    else:
        raise ValueError('Invalid ranks type')
    shape = A.shape
    cores = []
    C = A
    for i in range(len(shape) - 1):
        r_l, m, r_r = ranks[i], shape[i], ranks[i+1]
        unfolding = C.reshape(r_l * m, -1)
        U, S, Vt = torch.svd(unfolding)
        U, S, Vt = U[:, :r_r], S[:r_r], Vt[:r_r, :]
        cores.append(U.reshape(r_l, m, r_r))
        C = S.view(-1, 1) * Vt
    cores.append(C.reshape(ranks[-2], shape[-1], ranks[-1]))
    return cores
//  [python]
def tt_svd(A, ranks):
    warnings.warn('Use tntorch instead for a better approximation')
    if not torch.is_tensor(A):
        raise ValueError('Operand is not a tensor')
    if type(ranks) in (list, tuple):
        tt_ranks_max = get_tt_ranks(list(A.shape))
        if A.dim() + 1 != len(ranks) or ranks[0] != ranks[-1] != 1 or \
                any(a > b for a, b in zip(ranks, tt_ranks_max)):
            raise ValueError('Invalid ranks')
    elif type(ranks) is int:
        ranks = get_tt_ranks(A.shape, max_rank=ranks)
    else:
        raise ValueError('Invalid ranks type')
    shape = A.shape
    cores = []
    C = A
    for i in range(len(shape) - 1):
        r_l, m, r_r = ranks[i], shape[i], ranks[i+1]
        unfolding = C.reshape(r_l * m, -1)
        U, S, Vt = torch.linalg.svd(unfolding, full_matrices=False)
        U, S, Vt = U[:, :r_r], S[:r_r], Vt[:r_r, :]
        cores.append(U.reshape(r_l, m, r_r))
        C = S.view(-1, 1) * Vt
    cores.append(C.reshape(ranks[-2], shape[-1], ranks[-1]))
    return cores
//  [python]
def build_from_linear(cls, module: nn.Linear):
        with torch.no_grad():
            U, S, V_t = torch.svd(module.weight)

        result_module = cls(
            in_features=module.in_features,
            out_features=module.out_features,
            use_bias=module.bias is not None,
        )
        result_module.u.data = U
        result_module.v_t.data = V_t
        result_module.log_s.data = torch.log(S)
        if module.bias is not None:
            result_module.bias.data = module.bias
        return result_module
//  [python]
def build_from_linear(cls, module: nn.Linear):
        with torch.no_grad():
            U, S, V_t = torch.linalg.svd(module.weight, full_matrices=False)

        result_module = cls(
            in_features=module.in_features,
            out_features=module.out_features,
            use_bias=module.bias is not None,
        )
        result_module.u.data = U
        result_module.v_t.data = V_t
        result_module.log_s.data = torch.log(S)
        if module.bias is not None:
            result_module.bias.data = module.bias
        return result_module
//  [python]
def get_interpolated_path(cameras: PerspectiveCameras, n=50, method='circle', theta_offset_max=0.0, flip_up=False):
    '''
    Given a camera object containing a set of cameras, fit a circle and get 
    interpolated cameras

    Args:
        cameras (PyTorch3D Camera): input camera object
        n (int): length of cameras in new path
        method (str): 'circle'
        theta_offset_max (int): max camera jitter in radians

    Returns:
        path_cameras (PyTorch3D Camera): interpolated cameras
    '''
    device = cameras.device
    cameras = cameras.cpu()

    if method == 'circle':

        #@ https://meshlogic.github.io/posts/jupyter/curve-fitting/fitting-a-circle-to-cluster-of-3d-points/
        #@ Fit plane
        P = cameras.get_camera_center().cpu()
        P_mean = P.mean(axis=0)
        P_centered = P - P_mean
        U,s,V = torch.svd(P_centered)
        normal = V[2,:]
        if (normal*2 - P_mean).norm() < (normal - P_mean).norm():
            normal = - normal
        d = -torch.dot(P_mean, normal)  # d = -<p,n>    

        #@ Project pts to plane
        P_xy = rodrigues_rot(P_centered, normal, torch.tensor([0.0,0.0,1.0]))
        
        #@ Fit circle in 2D
        xc, yc, r = fit_circle_2d(P_xy[:,0], P_xy[:,1])
        t = torch.linspace(0, 2*math.pi, 100)
        xx = xc + r*torch.cos(t)
        yy = yc + r*torch.sin(t)

        #@ Project circle to 3D
        C = rodrigues_rot(torch.tensor([xc,yc,0.0]), torch.tensor([0.0,0.0,1.0]), normal) + P_mean
        C = C.flatten()

        #@ Get pts n 3D
        t = torch.linspace(0, 2*math.pi, n)
        u = P[0] - C
        new_camera_centers = generate_circle_by_vectors(t, C, r, normal, u)

        #@ OPTIONAL THETA OFFSET
        if theta_offset_max > 0.0:
            aug_theta = (torch.rand((new_camera_centers.shape[0])) * (2*theta_offset_max)) - theta_offset_max
            new_camera_centers = rodrigues_rot2(new_camera_centers, normal, aug_theta)

        #@ Get camera look at
        new_camera_look_at = get_nearest_centroid(cameras)
        factor = -1.0 if flip_up else 1.0

        #@ Get R T
        up_vec = -normal * factor
        R, T = look_at_view_transform(eye=new_camera_centers, at=new_camera_look_at.unsqueeze(0), up=up_vec.unsqueeze(0), device=cameras.device)
    else:
        raise NotImplementedError
    
    c = (cameras.principal_point).mean(dim=0, keepdim=True).expand(R.shape[0],-1)
    f = (cameras.focal_length).mean(dim=0, keepdim=True).expand(R.shape[0],-1)
    image_size = cameras.image_size[:1].expand(R.shape[0],-1)


    path_cameras = PerspectiveCameras(R=R,T=T,focal_length=f,principal_point=c,image_size=image_size, device=device)
    cameras = cameras.to(device)
    return path_cameras
//  [python]
def get_interpolated_path(cameras: PerspectiveCameras, n=50, method='circle', theta_offset_max=0.0, flip_up=False):
    '''
    Given a camera object containing a set of cameras, fit a circle and get 
    interpolated cameras

    Args:
        cameras (PyTorch3D Camera): input camera object
        n (int): length of cameras in new path
        method (str): 'circle'
        theta_offset_max (int): max camera jitter in radians

    Returns:
        path_cameras (PyTorch3D Camera): interpolated cameras
    '''
    device = cameras.device
    cameras = cameras.cpu()

    if method == 'circle':

        #@ https://meshlogic.github.io/posts/jupyter/curve-fitting/fitting-a-circle-to-cluster-of-3d-points/
        #@ Fit plane
        P = cameras.get_camera_center().cpu()
        P_mean = P.mean(axis=0)
        P_centered = P - P_mean
        U,s,V = torch.linalg.svd(P_centered)
        normal = V[2,:]
        if (normal*2 - P_mean).norm() < (normal - P_mean).norm():
            normal = - normal
        d = -torch.dot(P_mean, normal)  # d = -<p,n>    

        #@ Project pts to plane
        P_xy = rodrigues_rot(P_centered, normal, torch.tensor([0.0,0.0,1.0]))
        
        #@ Fit circle in 2D
        xc, yc, r = fit_circle_2d(P_xy[:,0], P_xy[:,1])
        t = torch.linspace(0, 2*math.pi, 100)
        xx = xc + r*torch.cos(t)
        yy = yc + r*torch.sin(t)

        #@ Project circle to 3D
        C = rodrigues_rot(torch.tensor([xc,yc,0.0]), torch.tensor([0.0,0.0,1.0]), normal) + P_mean
        C = C.flatten()

        #@ Get pts n 3D
        t = torch.linspace(0, 2*math.pi, n)
        u = P[0] - C
        new_camera_centers = generate_circle_by_vectors(t, C, r, normal, u)

        #@ OPTIONAL THETA OFFSET
        if theta_offset_max > 0.0:
            aug_theta = (torch.rand((new_camera_centers.shape[0])) * (2*theta_offset_max)) - theta_offset_max
            new_camera_centers = rodrigues_rot2(new_camera_centers, normal, aug_theta)

        #@ Get camera look at
        new_camera_look_at = get_nearest_centroid(cameras)
        factor = -1.0 if flip_up else 1.0

        #@ Get R T
        up_vec = -normal * factor
        R, T = look_at_view_transform(eye=new_camera_centers, at=new_camera_look_at.unsqueeze(0), up=up_vec.unsqueeze(0), device=cameras.device)
    else:
        raise NotImplementedError
    
    c = (cameras.principal_point).mean(dim=0, keepdim=True).expand(R.shape[0],-1)
    f = (cameras.focal_length).mean(dim=0, keepdim=True).expand(R.shape[0],-1)
    image_size = cameras.image_size[:1].expand(R.shape[0],-1)


    path_cameras = PerspectiveCameras(R=R,T=T,focal_length=f,principal_point=c,image_size=image_size, device=device)
    cameras = cameras.to(device)
    return path_cameras
//  [python]
def compute_optimal_alignment(A, B):
    """
    Compute the optimal scale s, rotation R, and translation t that minimizes:
    || A - (s * B @ R + T) || ^ 2

    Reference: Umeyama (TPAMI 91)

    Args:
        A (torch.Tensor): (N, 3).
        B (torch.Tensor): (N, 3).

    Returns:
        s (float): scale.
        R (torch.Tensor): rotation matrix (3, 3).
        t (torch.Tensor): translation (3,).
    """
    A_bar = A.mean(0)
    B_bar = B.mean(0)
    # normally with R @ B, this would be A @ B.T
    H = (B - B_bar).T @ (A - A_bar)
    U, S, Vh = torch.svd(H)
    s = torch.det(U @ Vh)
    S_prime = torch.diag(torch.tensor([1, 1, torch.sign(s)], device=A.device))
    variance = torch.sum((B - B_bar) ** 2)
    scale = 1 / variance * torch.trace(torch.diag(S) @ S_prime)
    R = U @ S_prime @ Vh
    t = A_bar - scale * B_bar @ R

    A_hat = scale * B @ R + t
    return A_hat, scale, R, t
//  [python]
def compute_optimal_alignment(A, B):
    """
    Compute the optimal scale s, rotation R, and translation t that minimizes:
    || A - (s * B @ R + T) || ^ 2

    Reference: Umeyama (TPAMI 91)

    Args:
        A (torch.Tensor): (N, 3).
        B (torch.Tensor): (N, 3).

    Returns:
        s (float): scale.
        R (torch.Tensor): rotation matrix (3, 3).
        t (torch.Tensor): translation (3,).
    """
    A_bar = A.mean(0)
    B_bar = B.mean(0)
    # normally with R @ B, this would be A @ B.T
    H = (B - B_bar).T @ (A - A_bar)
    U, S, Vh = torch.linalg.svd(H, full_matrices=True)
    s = torch.linalg.det(U @ Vh)
    S_prime = torch.diag(torch.tensor([1, 1, torch.sign(s)], device=A.device))
    variance = torch.sum((B - B_bar) ** 2)
    scale = 1 / variance * torch.trace(torch.diag(S) @ S_prime)
    R = U @ S_prime @ Vh
    t = A_bar - scale * B_bar @ R

    A_hat = scale * B @ R + t
    return A_hat, scale, R, t
//  [python]
def _torch_svd_cast(self, input):
        """Helper function to make torch.svd work with other than fp32/64.

        The function torch.svd is only implemented for fp32/64 which makes
        impossible to be used by fp16 or others. What this function does, is cast
        input data type to fp32, apply torch.svd, and cast back to the input dtype.

        NOTE: in torch 1.8.1 this function is recommended to use as torch.linalg.svd
        """
        # if not isinstance(input, torch.Tensor):
        #    raise AssertionError(f"Input must be torch.Tensor. Got: {type(input)}.")
        dtype = input.dtype
        if dtype not in (torch.float32, torch.float64):
            dtype = torch.float32

        out1, out2, out3H = torch.svd(input.to(dtype))
        #if torch_version_ge(1, 11):
        #    out3 = out3H.mH
        #else:
        out3 = out3H.transpose(-1, -2)
        return (out1.to(input.dtype), out2.to(input.dtype), out3.to(input.dtype))
//  [python]
def _torch_svd_cast(self, input):
        """Helper function to make torch.svd work with other than fp32/64.

        The function torch.svd is only implemented for fp32/64 which makes
        impossible to be used by fp16 or others. What this function does, is cast
        input data type to fp32, apply torch.svd, and cast back to the input dtype.

        NOTE: in torch 1.8.1 this function is recommended to use as torch.linalg.svd
        """
        # if not isinstance(input, torch.Tensor):
        #    raise AssertionError(f"Input must be torch.Tensor. Got: {type(input)}.")
        dtype = input.dtype
        if dtype not in (torch.float32, torch.float64):
            dtype = torch.float32

        out1, out2, out3H = torch.linalg.svd(input.to(dtype))
        #if torch_version_ge(1, 11):
        #    out3 = out3H.mH
        #else:
        out3 = out3H.transpose(-1, -2)
        return (out1.to(input.dtype), out2.to(input.dtype), out3.to(input.dtype))
//  [python]
def get_rmsd(a, b, eps=1e-6):
    '''
    align crds b to a : always use all alphas
    expexted tensor of shape (L,3)
    jake's torch adapted version
    '''
    assert a.shape == b.shape, 'make sure tensors are the same size'
    L = a.shape[0]
    assert a.shape == torch.Size([L,3]), 'make sure tensors are in format [L,3]'

    # center to CA centroid
    a = a - a.mean(dim=0)
    b = b - b.mean(dim=0)

    # Computation of the covariance matrix
    C = torch.einsum('kj,ji->ki', torch.transpose(b.type(torch.float32),0,1), a.type(torch.float32))

    # Compute optimal rotation matrix using SVD
    U, S, V = torch.svd(C)

    # get sign to ensure right-handedness
    d = torch.ones([3,3])
    d[:,-1] = torch.sign(torch.det(U)*torch.det(V))

    # Rotation matrix U
    U = torch.einsum('kj,ji->ki',(d*U),V)

    # Rotate xyz_hal
    rP = torch.einsum('kj,ji->ki',b.type(torch.float32),U.type(torch.float32))

    L = rP.shape[0]
    rmsd = torch.sqrt(torch.sum((rP-a)*(rP-a), axis=(0,1)) / L + eps)

    return rmsd, U
//  [python]
def get_rmsd(a, b, eps=1e-6):
    '''
    align crds b to a : always use all alphas
    expexted tensor of shape (L,3)
    jake's torch adapted version
    '''
    assert a.shape == b.shape, 'make sure tensors are the same size'
    L = a.shape[0]
    assert a.shape == torch.Size([L,3]), 'make sure tensors are in format [L,3]'

    # center to CA centroid
    a = a - a.mean(dim=0)
    b = b - b.mean(dim=0)

    # Computation of the covariance matrix
    C = torch.einsum('kj,ji->ki', torch.transpose(b.type(torch.float32),0,1), a.type(torch.float32))

    # Compute optimal rotation matrix using SVD
    V, S, W = torch.linalg.svd(C)

    # get sign to ensure right-handedness
    d = torch.ones([3,3])
    d[:,-1] = torch.sign(torch.linalg.det(V)*torch.linalg.det(W))

    # Rotation matrix U
    U = torch.einsum('kj,ji->ki',(d*V),W)

    # Rotate xyz_hal
    rP = torch.einsum('kj,ji->ki',b.type(torch.float32),U.type(torch.float32))

    L = rP.shape[0]
    rmsd = torch.sqrt(torch.sum((rP-a)*(rP-a), axis=(0,1)) / L + eps)

    return rmsd, U
//  [python]
def forward(ctx, input, epsilon=None):
        p = input

        try:
            _, eig_diag, eig_vec_transposed = torch.svd(p, some=False)
            eig_vec = eig_vec_transposed.transpose(-2, -1)  # Equivalent result to the commented line

        except LinAlgError as error:
            save_error_tensor(p, error, "torch.svd", Eigen_decomposition.forward, Eigen_decomposition)
            raise error

        dtype = eig_diag.dtype
        if epsilon is None:
            epsilon = torch.finfo(dtype).eps
        else:
            epsilon = max(epsilon, torch.finfo(dtype).eps)

        eig_diag[eig_diag <= epsilon] = epsilon  # Zero-out eigenvalues smaller than epsilon
        ctx.save_for_backward(eig_vec, eig_diag)
        return eig_vec, eig_diag
//  [python]
def forward(ctx, input, epsilon=None):
        p = input

        try:
            _, eig_diag, eig_vec_transposed = torch.linalg.svd(p, full_matrices=False)
            eig_vec = eig_vec_transposed.transpose(-2, -1)  # Equivalent result to the commented line

        except LinAlgError as error:
            save_error_tensor(p, error, "torch.linalg.svd", Eigen_decomposition.forward, Eigen_decomposition)
            raise error

        dtype = eig_diag.dtype
        if epsilon is None:
            epsilon = torch.finfo(dtype).eps
        else:
            epsilon = max(epsilon, torch.finfo(dtype).eps)

        eig_diag[eig_diag <= epsilon] = epsilon  # Zero-out eigenvalues smaller than epsilon
        ctx.save_for_backward(eig_vec, eig_diag)
        return eig_vec, eig_diag
//  [python]
def DARE_GRAM_LOSS(H1, H2):    
    b,p = H1.shape

    A = torch.cat((torch.ones(b,1).to(device), H1), 1)
    B = torch.cat((torch.ones(b,1).to(device), H2), 1)

    cov_A = (A.t()@A)
    cov_B = (B.t()@B) 

    U_A,L_A,V_A = torch.svd(cov_A)
    U_B,L_B,V_B = torch.svd(cov_B)
    
    eigen_A = torch.cumsum(L_A.detach(), dim=0)/L_A.sum()
    eigen_B = torch.cumsum(L_B.detach(), dim=0)/L_B.sum()

    if(eigen_A[1]>args.treshold):
        T = eigen_A[1].detach()
    else:
        T = args.treshold
        
    index_A = torch.argwhere(eigen_A.detach()<=T)[-1]

    if(eigen_B[1]>args.treshold):
        T = eigen_B[1].detach()
    else:
        T = args.treshold

    index_B = torch.argwhere(eigen_B.detach()<=T)[-1]
    
    k = max(index_A, index_B)[0]

    A = torch.linalg.pinv(cov_A ,rtol = (L_A[k]/L_A[0]).detach())
    B = torch.linalg.pinv(cov_B ,rtol = (L_B[k]/L_B[0]).detach())
    
    cos_sim = nn.CosineSimilarity(dim=0,eps=1e-6)
    cos = torch.dist(torch.ones((p+1)).to(device),(cos_sim(A,B)),p=1)/(p+1)
    return args.tradeoff_angle*(cos) + args.tradeoff_scale*torch.dist((L_A[:k]),(L_B[:k]))/k
//  [python]
def DARE_GRAM_LOSS(H1, H2):    
    b,p = H1.shape

    A = torch.cat((torch.ones(b,1).to(device), H1), 1)
    B = torch.cat((torch.ones(b,1).to(device), H2), 1)

    cov_A = (A.t()@A)
    cov_B = (B.t()@B) 

    _,L_A,_ = torch.linalg.svd(cov_A)
    _,L_B,_ = torch.linalg.svd(cov_B)
    
    eigen_A = torch.cumsum(L_A.detach(), dim=0)/L_A.sum()
    eigen_B = torch.cumsum(L_B.detach(), dim=0)/L_B.sum()

    if(eigen_A[1]>args.treshold):
        T = eigen_A[1].detach()
    else:
        T = args.treshold
        
    index_A = torch.argwhere(eigen_A.detach()<=T)[-1]

    if(eigen_B[1]>args.treshold):
        T = eigen_B[1].detach()
    else:
        T = args.treshold

    index_B = torch.argwhere(eigen_B.detach()<=T)[-1]
    
    k = max(index_A, index_B)[0]

    A = torch.linalg.pinv(cov_A ,rtol = (L_A[k]/L_A[0]).detach())
    B = torch.linalg.pinv(cov_B ,rtol = (L_B[k]/L_B[0]).detach())
    
    cos_sim = nn.CosineSimilarity(dim=0,eps=1e-6)
    cos = torch.dist(torch.ones((p+1)).to(device),(cos_sim(A,B)),p=1)/(p+1)
    return args.tradeoff_angle*(cos) + args.tradeoff_scale*torch.dist((L_A[:k]),(L_B[:k]))/k
//  [python]
def ssa(y, dim) -> tuple:
    """
    Singular Spectrum Analysis decomposition for a time series

    Example:
    -------
    >>> import numpy as np
    >>>
    >>> x = np.linspace(0, 5, 1000)
    >>> y = 2*x + 2*np.sin(5*x) + 0.5*np.random.randn(1000)
    >>> pc, s, v = ssa(y, 15)

    :param y: time series (array)
    :param dim: the embedding dimension
    :return: (pc, s, v) where
             pc is the matrix with the principal components of y
             s is the vector of the singular values of y given dim
             v is the matrix of the singular vectors of y given dim

    """
    with torch.no_grad():
        n = len(y)
        t = n - (dim - 1)

        t1 = time.time()
        yy = linalg.hankel(y, np.zeros(dim))
        yy = yy[:-dim + 1, :] / np.sqrt(t)
        t2 = time.time()

        # here we use gesvd driver (as in Matlab)
        _, s, v = linalg.svd(yy, full_matrices=False, lapack_driver='gesvd')
        # yy = torch.from_numpy(yy).to(device)
        # yy = torch.from_numpy(yy)
        # _, s, v = torch.linalg.svd(yy, full_matrices=False)
        # s = s.cpu().numpy()
        # v = v.cpu().numpy()
        t3 = time.time()

        # find principal components
        # vt = np.matrix(v).T
        # pc = np.matrix(yy) * vt
        # s = s.to(device)
        # v = v.to(device)
        vt = v.transpose(0, 1)
        # pc = torch.matmul(yy, vt)
        pc = np.matmul(yy, vt)
        t4 = time.time()
        print(t2-t1, t3-t2, t4-t3)

        # return np.asarray(pc), s, np.asarray(vt)
        # return pc.cpu().numpy(), s.cpu().numpy(), vt.cpu().numpy()
        return pc, s, vt
//  [python]
def ssa(y, dim) -> tuple:
    """
    Singular Spectrum Analysis decomposition for a time series

    Example:
    -------
    >>> import numpy as np
    >>>
    >>> x = np.linspace(0, 5, 1000)
    >>> y = 2*x + 2*np.sin(5*x) + 0.5*np.random.randn(1000)
    >>> pc, s, v = ssa(y, 15)

    :param y: time series (array)
    :param dim: the embedding dimension
    :return: (pc, s, v) where
             pc is the matrix with the principal components of y
             s is the vector of the singular values of y given dim
             v is the matrix of the singular vectors of y given dim

    """
    with torch.no_grad():
        n = len(y)
        t = n - (dim - 1)

        t1 = time.time()
        yy = linalg.hankel(y, np.zeros(dim))
        yy = yy[:-dim + 1, :] / np.sqrt(t)
        t2 = time.time()

        # here we use gesvd driver (as in Matlab)
        # _, s, v = linalg.svd(yy, full_matrices=False, lapack_driver='gesvd')
        # yy = torch.from_numpy(yy).to(device)
        yy = torch.from_numpy(yy)
        _, s, v = torch.linalg.svd(yy, full_matrices=False)
        # s = s.cpu().numpy()
        # v = v.cpu().numpy()
        t3 = time.time()

        # find principal components
        # vt = np.matrix(v).T
        # pc = np.matrix(yy) * vt
        # s = s.to(device)
        # v = v.to(device)
        vt = v.transpose(0, 1)
        # pc = torch.matmul(yy, vt)
        pc = torch.matmul(yy, vt)
        t4 = time.time()
        print(t2-t1, t3-t2, t4-t3)

        # return np.asarray(pc), s, np.asarray(vt)
        # return pc.cpu().numpy(), s.cpu().numpy(), vt.cpu().numpy()
        return pc, s, vt
//  [python]
def fit_plane(points):
    """
    :param points (*, N, 3)
    returns (*, 3) plane parameters (returns in normal * offset format)
    """
    *dims, N, D = points.shape
    mean = points.mean(dim=-2, keepdim=True)
    # (*, N, D), (*, D), (*, D, D)
    U, S, Vh = torch.svd(points - mean)
    normal = Vh[..., -1, :]  # (*, D)
    offset = torch.einsum("...ij,...j->...i", points, normal)  # (*, N)
    offset = offset.mean(dim=-1, keepdim=True)
    return torch.cat([normal, offset], dim=-1)
//  [python]
def fit_plane(points):
    """
    :param points (*, N, 3)
    returns (*, 3) plane parameters (returns in normal * offset format)
    """
    *dims, N, D = points.shape
    mean = points.mean(dim=-2, keepdim=True)
    # (*, N, D), (*, D), (*, D, D)
    U, S, Vh = torch.linalg.svd(points - mean)
    normal = Vh[..., -1, :]  # (*, D)
    offset = torch.einsum("...ij,...j->...i", points, normal)  # (*, N)
    offset = offset.mean(dim=-1, keepdim=True)
    return torch.cat([normal, offset], dim=-1)
//  [python]
def forward(
        self,
        x: torch.Tensor,
        features: int,
        num_itt: int = 10,
        lam: float = 0.01,
        normalize: bool = False,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Denoise an image `x` using the HyRes algorithm.

        Parameters
        ----------
        x: torch.Tensor
            input image
        features: int
            Number of features to extract, this could be selected equal to the
            number of classes of interests in the scene
        num_itt: int, optional
            Number of iterations; 200 iterations are default value
        lam: float, optional
            Tuning parameter; Default is 0.01 for normalized HSI
            the value passed to the denoising algorithm is 1/lam
        normalize: bool, optional
            flag indicating if the input should be normalized

        Returns
        -------
        denoised_image : torch.Tensor
            Hyperspectral features extracted (3D matrix)
        fe : torch.Tensor
            Extracted features
        """
        nr1, nc1, p1 = x.shape

        if normalize:
            x, consts = utils.normalize(x, by_band=True)

        x_2d = x.reshape((nr1 * nc1, p1))
        U, S, Vh = torch.svd(x_2d)
        v1 = Vh.T
        v = v1[:, :features]
        fe = torch.zeros((nr1, nc1, features), dtype=x.dtype, device=x.device)

        for fi in range(num_itt):
            c1 = x_2d @ v[:, :features]
            pc = c1.reshape((nr1, nc1, features))

            fe = denoise_tv_bregman(image=pc.cpu().numpy(), weight=1 / lam, eps=0.1, isotropic=True)
            fe = torch.tensor(fe, dtype=x.dtype, device=x.device)

            fe_reshape = fe.reshape((nr1 * nc1, features))
            m = x_2d.T @ fe_reshape
            U, S, Vh = torch.svd(m)
            v = U @ Vh

        denoised_image = (fe_reshape @ v.T).reshape((nr1, nc1, p1))
        if normalize:
            denoised_image = utils.undo_normalize(denoised_image, **consts, by_band=True)
        return denoised_image, fe
//  [python]
def forward(
        self,
        x: torch.Tensor,
        features: int,
        num_itt: int = 10,
        lam: float = 0.01,
        normalize: bool = False,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Denoise an image `x` using the HyRes algorithm.

        Parameters
        ----------
        x: torch.Tensor
            input image
        features: int
            Number of features to extract, this could be selected equal to the
            number of classes of interests in the scene
        num_itt: int, optional
            Number of iterations; 200 iterations are default value
        lam: float, optional
            Tuning parameter; Default is 0.01 for normalized HSI
            the value passed to the denoising algorithm is 1/lam
        normalize: bool, optional
            flag indicating if the input should be normalized

        Returns
        -------
        denoised_image : torch.Tensor
            Hyperspectral features extracted (3D matrix)
        fe : torch.Tensor
            Extracted features
        """
        nr1, nc1, p1 = x.shape

        if normalize:
            x, consts = utils.normalize(x, by_band=True)

        x_2d = x.reshape((nr1 * nc1, p1))
        _, _, vh = torch.linalg.svd(x_2d, full_matrices=False)
        v1 = vh.T
        v = v1[:, :features]
        fe = torch.zeros((nr1, nc1, features), dtype=x.dtype, device=x.device)

        for fi in range(num_itt):
            c1 = x_2d @ v[:, :features]
            pc = c1.reshape((nr1, nc1, features))

            fe = denoise_tv_bregman(image=pc.cpu().numpy(), weight=1 / lam, eps=0.1, isotropic=True)
            fe = torch.tensor(fe, dtype=x.dtype, device=x.device)

            fe_reshape = fe.reshape((nr1 * nc1, features))
            m = x_2d.T @ fe_reshape
            c, _, gh = torch.linalg.svd(m, full_matrices=False)
            v = c @ gh

        denoised_image = (fe_reshape @ v.T).reshape((nr1, nc1, p1))
        if normalize:
            denoised_image = utils.undo_normalize(denoised_image, **consts, by_band=True)
        return denoised_image, fe
//  [python]
def torch_svd(x):
    u, s, vh = torch.svd(x)
    
    print("svd: ", x.size(), u.size(), s.size(), vh.size())
    h = u[:, :s.size(0)] * torch.sqrt(s)
    return h
//  [python]
def torch_svd(x):
    u, s, vh = torch.linalg.svd(x)
    
    print("svd: ", x.size(), u.size(), s.size(), vh.size())
    h = u[:, :s.size(0)] * torch.sqrt(s)
    return h
//  [python]
def modalsvd(self, inputs, n):
        """
        Method that performs the standard SVD of an unfolded matrix defined from an input tensor
        along a given dimension

        :param torch.Tensor inputs: the input tensor
        :param int n: the index of the unfolding matrix being decomposed
        :return: three SVD matrices of the n-th unfolding of tensor inputs
        :rtype: torch.Tensors
        """
        return torch.svd(self.unfolding(inputs, n), some=True)
//  [python]
def modalsvd(self, inputs, n):
        """
        Method that performs the standard SVD of an unfolded matrix defined from an input tensor
        along a given dimension

        :param torch.Tensor inputs: the input tensor
        :param int n: the index of the unfolding matrix being decomposed
        :return: three SVD matrices of the n-th unfolding of tensor inputs
        :rtype: torch.Tensors
        """
        return torch.linalg.svd(self.unfolding(inputs, n), full_matrices=True)
//  [python]
def orth_procrustes(x, y):
    x = normalize(x)
    y = normalize(y)
    xty = torch.einsum('bvi,bvj->bij', y, x)
    u, s, v = torch.svd(xty)
    r = torch.einsum('bij,bkj->bik', u, v)
    return r
//  [python]
def orth_procrustes(x, y):
    x = normalize(x)
    y = normalize(y)
    xty = torch.einsum('bvi,bvj->bij', y, x)
    s, u, v = torch.linalg.svd(xty)
    r = torch.einsum('bij,bkj->bik', u, v)
    return r
//  [python]
def test_linalg(device="cpu") -> None:
    print(f"Testing smoke_test_linalg on {device}")
    A = torch.randn(5, 3, device=device)
    U, S, Vh = torch.svd(A, some=False)
    assert U.shape == A.shape and S.shape == torch.Size([3]) and Vh.shape == torch.Size([3, 3])
    torch.dist(A, U @ torch.diag(S) @ Vh)

    U, S, Vh = torch.svd(A)
    assert U.shape == torch.Size([5, 5]) and S.shape == torch.Size([3]) and Vh.shape == torch.Size([3, 3])
    torch.dist(A, U[:, :3] @ torch.diag(S) @ Vh)

    A = torch.randn(7, 5, 3, device=device)
    U, S, Vh = torch.svd(A, some=False)
    torch.dist(A, U @ torch.diag_embed(S) @ Vh)

    if device == "cuda":
        supported_dtypes = [torch.float32, torch.float64]
        for dtype in supported_dtypes:
            print(f"Testing smoke_test_linalg with cuda for {dtype}")
            A = torch.randn(20, 16, 50, 100, device=device, dtype=dtype)
            torch.svd(A)
//  [python]
def test_linalg(device="cpu") -> None:
    print(f"Testing smoke_test_linalg on {device}")
    A = torch.randn(5, 3, device=device)
    U, S, Vh = torch.linalg.svd(A, full_matrices=False)
    assert U.shape == A.shape and S.shape == torch.Size([3]) and Vh.shape == torch.Size([3, 3])
    torch.dist(A, U @ torch.diag(S) @ Vh)

    U, S, Vh = torch.linalg.svd(A)
    assert U.shape == torch.Size([5, 5]) and S.shape == torch.Size([3]) and Vh.shape == torch.Size([3, 3])
    torch.dist(A, U[:, :3] @ torch.diag(S) @ Vh)

    A = torch.randn(7, 5, 3, device=device)
    U, S, Vh = torch.linalg.svd(A, full_matrices=False)
    torch.dist(A, U @ torch.diag_embed(S) @ Vh)

    if device == "cuda":
        supported_dtypes = [torch.float32, torch.float64]
        for dtype in supported_dtypes:
            print(f"Testing smoke_test_linalg with cuda for {dtype}")
            A = torch.randn(20, 16, 50, 100, device=device, dtype=dtype)
            torch.linalg.svd(A)
//  [python]
def find_rigid_alignment(A, B):
    """
    See: https://en.wikipedia.org/wiki/Kabsch_algorithm
    2-D or 3-D registration with known correspondences.
    Registration occurs in the zero centered coordinate system, and then
    must be transported back.
        Args:
        -    A: Torch tensor of shape (N,D) -- Point Cloud to Align (source)
        -    B: Torch tensor of shape (N,D) -- Reference Point Cloud (target)
        Returns:
        -    R: optimal rotation
        -    t: optimal translation
    Test on rotation + translation and on rotation + translation + reflection
    """
    a_mean = A.mean(axis=0)
    b_mean = B.mean(axis=0)
    A_c = A - a_mean
    B_c = B - b_mean
    # Covariance matrix
    H = A_c.T.mm(B_c)
    U, S, Vt = torch.svd(H)
    # Rotation matrix
    R = Vt.T.mm(U.T)
    
    # special reflection case
    if torch.det(R) < 0:
        # print("det(R) < R, reflection detected!, correcting for it ...")
        SS = torch.diag(torch.tensor([1.,1.,-1.], device=R.device))
        R = (Vt.T @ SS) @ U.T
    
    # Translation vector
    t = b_mean[None, :] - R.mm(a_mean[None, :].T).T
    t = t.T
    return R, t.squeeze()
//  [python]
def find_rigid_alignment(A, B):
    """
    See: https://en.wikipedia.org/wiki/Kabsch_algorithm
    2-D or 3-D registration with known correspondences.
    Registration occurs in the zero centered coordinate system, and then
    must be transported back.
        Args:
        -    A: Torch tensor of shape (N,D) -- Point Cloud to Align (source)
        -    B: Torch tensor of shape (N,D) -- Reference Point Cloud (target)
        Returns:
        -    R: optimal rotation
        -    t: optimal translation
    Test on rotation + translation and on rotation + translation + reflection
    """
    a_mean = A.mean(axis=0)
    b_mean = B.mean(axis=0)
    A_c = A - a_mean
    B_c = B - b_mean
    # Covariance matrix
    H = A_c.T.mm(B_c)
    U, S, Vt = torch.linalg.svd(H)
    # Rotation matrix
    R = Vt.T.mm(U.T)
    
    # special reflection case
    if torch.linalg.det(R) < 0:
        # print("det(R) < R, reflection detected!, correcting for it ...")
        SS = torch.diag(torch.tensor([1.,1.,-1.], device=R.device))
        R = (Vt.T @ SS) @ U.T
    
    # Translation vector
    t = b_mean[None, :] - R.mm(a_mean[None, :].T).T
    t = t.T
    return R, t.squeeze()
//  [python]
def _get_nearest_centroid(cameras: PerspectiveCameras):
    '''
    Given PyTorch3D cameras, find the nearest point along their principal ray
    '''

    #@ GET CAMERA CENTERS AND DIRECTIONS
    camera_centers = cameras.get_camera_center()

    c_mean = (cameras.principal_point).mean(dim=0)
    xy_grid = c_mean.unsqueeze(0).unsqueeze(0)
    ray_vis = _xy_to_ray_bundle(cameras, xy_grid.expand(len(cameras),-1,-1), 1.0, 15.0, 20, True)
    camera_directions = ray_vis.directions

    #@ CONSTRUCT MATRICIES
    A = torch.zeros((3*len(cameras)), len(cameras)+3)
    b = torch.zeros((3*len(cameras), 1))
    A[:,:3] = torch.eye(3).repeat(len(cameras),1)
    for ci in range(len(camera_directions)):
        A[3*ci:3*ci+3, ci+3] = -camera_directions[ci]
        b[3*ci:3*ci+3, 0] = camera_centers[ci]
    #' A (3*N, 3*N+3)   b (3*N, 1)

    #@ SVD
    U, s, VT = torch.svd(A)
    Sinv = torch.diag(1/s)
    if len(s) < 3*len(cameras):
        Sinv = torch.cat((Sinv, torch.zeros((Sinv.shape[0], 3*len(cameras) - Sinv.shape[1]), device=Sinv.device)), dim=1)
    x = torch.matmul(VT.T, torch.matmul(Sinv,torch.matmul(U.T, b)))
    
    centroid = x[:3,0]
    return centroid
//  [python]
def _get_nearest_centroid(cameras: PerspectiveCameras):
    '''
    Given PyTorch3D cameras, find the nearest point along their principal ray
    '''

    #@ GET CAMERA CENTERS AND DIRECTIONS
    camera_centers = cameras.get_camera_center()

    c_mean = (cameras.principal_point).mean(dim=0)
    xy_grid = c_mean.unsqueeze(0).unsqueeze(0)
    ray_vis = _xy_to_ray_bundle(cameras, xy_grid.expand(len(cameras),-1,-1), 1.0, 15.0, 20, True)
    camera_directions = ray_vis.directions

    #@ CONSTRUCT MATRICIES
    A = torch.zeros((3*len(cameras)), len(cameras)+3)
    b = torch.zeros((3*len(cameras), 1))
    A[:,:3] = torch.eye(3).repeat(len(cameras),1)
    for ci in range(len(camera_directions)):
        A[3*ci:3*ci+3, ci+3] = -camera_directions[ci]
        b[3*ci:3*ci+3, 0] = camera_centers[ci]
    #' A (3*N, 3*N+3)   b (3*N, 1)

    #@ SVD
    U, s, VT = torch.linalg.svd(A)
    Sinv = torch.diag(1/s)
    if len(s) < 3*len(cameras):
        Sinv = torch.cat((Sinv, torch.zeros((Sinv.shape[0], 3*len(cameras) - Sinv.shape[1]), device=Sinv.device)), dim=1)
    x = torch.matmul(VT.T, torch.matmul(Sinv,torch.matmul(U.T, b)))
    
    centroid = x[:3,0]
    return centroid
//  [python]
def pca_by_group(point_xyz, group_ids):
    """
    Args:
        point_xyz [N, 3]
        group_ids [N] in range [G]
    Returns:
        eigvals [G, 3]
        eigvecs [G, 3, 3]
    """
    num_groups = group_ids.max().item() + 1

    group_center = scatter(point_xyz, group_ids, dim=0,
                           dim_size=num_groups, reduce='mean')
    
    point_d = point_xyz - group_center[group_ids] # [N, 3]

    point_ddT = point_d.unsqueeze(-1) @ point_d.unsqueeze(-2) # [N, 3, 3]
    group_cov = scatter(point_ddT, group_ids, dim=0,
                        dim_size=num_groups, reduce='mean') # [G, 3, 3]
    group_cov = group_cov + torch.diag_embed(torch.tensor([1, 1, 1]).float()).to(cov).repeat(num_groups, 1, 1) * self.theta1
    U, S, V = torch.linalg.svd(group_cov)
    eigvals = S
    eigvecs = U @ V.transpose(-2, -1)
    return eigvals, eigvecs
//  [python]
def pca_by_group(point_xyz, group_ids):
    """
    Args:
        point_xyz [N, 3]
        group_ids [N] in range [G]
    Returns:
        eigvals [G, 3]
        eigvecs [G, 3, 3]
    """
    num_groups = group_ids.max().item() + 1

    group_center = scatter(point_xyz, group_ids, dim=0,
                           dim_size=num_groups, reduce='mean')
    
    point_d = point_xyz - group_center[group_ids] # [N, 3]

    point_ddT = point_d.unsqueeze(-1) @ point_d.unsqueeze(-2) # [N, 3, 3]
    group_cov = scatter(point_ddT, group_ids, dim=0,
                        dim_size=num_groups, reduce='mean') # [G, 3, 3]
    group_cov = group_cov + torch.diag_embed(torch.tensor([1, 1, 1]).float()).to(cov).repeat(num_groups, 1, 1) * self.theta1
    eigvecs, eigvals = torch.linalg.svd(group_cov)
    return eigvals, eigvecs
//  [python]
def __init__(self, skeletons, asymmetric=False):
		# representation
		ridxs = TOPOLOGY.roots(indices=True)
		self._jidxs = [i for i in range(len(TOPOLOGY)) if i not in ridxs]
		
		# handle LR symmetry
		self._asymmetric = bool(asymmetric)
		sorted_joints = sorted(enumerate(TOPOLOGY), key=lambda t: t[1])
		self._left_jidxs = [i for i, j in sorted_joints if "left" in j]
		self._right_jidxs = [i for i, j in sorted_joints if "right" in j]
		
		# Compute bone lengths and unit offsets
		offsets = Positions.to_local(skeletons, TOPOLOGY)[..., self._jidxs, :]		# S x J x 3
		lengths = offsets.norm(p=2, dim=-1, keepdim=True)							# S x J x 1
		self._unit_offsets = torch.zeros(len(TOPOLOGY), 3)							# J x 3
		self._unit_offsets[self._jidxs, :] = offsets.where(lengths <= 1e-10, offsets / lengths).mean(dim=0)
		
		# Svd Basis
		svd_samples = lengths.squeeze(-1)
		self._mean = svd_samples.mean(dim=0)										# C
		svd_samples -= self._mean
		U, S, V = torch.svd(svd_samples)										# C x C
		basis = svd_samples * V.unsqueeze(1)										# C x N x C
		self._basis = basis.sum(dim=2).std(dim=1, keepdim=True) * V
//  [python]
def __init__(self, skeletons, asymmetric=False):
		# representation
		ridxs = TOPOLOGY.roots(indices=True)
		self._jidxs = [i for i in range(len(TOPOLOGY)) if i not in ridxs]
		
		# handle LR symmetry
		self._asymmetric = bool(asymmetric)
		sorted_joints = sorted(enumerate(TOPOLOGY), key=lambda t: t[1])
		self._left_jidxs = [i for i, j in sorted_joints if "left" in j]
		self._right_jidxs = [i for i, j in sorted_joints if "right" in j]
		
		# Compute bone lengths and unit offsets
		offsets = Positions.to_local(skeletons, TOPOLOGY)[..., self._jidxs, :]		# S x J x 3
		lengths = offsets.norm(p=2, dim=-1, keepdim=True)							# S x J x 1
		self._unit_offsets = torch.zeros(len(TOPOLOGY), 3)							# J x 3
		self._unit_offsets[self._jidxs, :] = offsets.where(lengths <= 1e-10, offsets / lengths).mean(dim=0)
		
		# Svd Basis
		svd_samples = lengths.squeeze(-1)
		self._mean = svd_samples.mean(dim=0)										# C
		svd_samples -= self._mean
		V = torch.linalg.svd(svd_samples)[2]										# C x C
		basis = svd_samples * V.unsqueeze(1)										# C x N x C
		self._basis = basis.sum(dim=2).std(dim=1, keepdim=True) * V
//  [python]
def pca_for(self, x):
        batch_size, feature_dim = x.size()
        f_cov = (torch.mm(x.transpose(0, 1), x) / (batch_size - 1)).float()  # d * d
        eye = torch.eye(feature_dim).float().to(f_cov.device)
        # f_cov = torch.mm(x.transpose(0, 1), x).float()  # d * d
        U, S, V = torch.svd(0.99 * f_cov + 0.01 * eye)
        diag = torch.diag(1.0 / torch.sqrt(S + 1e-5))
        rotate_mtx = torch.mm(diag, U.transpose(0, 1)).detach()  # d * d
        return torch.mm(rotate_mtx, x.transpose(0, 1)).transpose(0, 1)
//  [python]
def pca_for(self, x):
        batch_size, feature_dim = x.size()
        f_cov = (torch.mm(x.transpose(0, 1), x) / (batch_size - 1)).float()  # d * d
        eye = torch.eye(feature_dim).float().to(f_cov.device)
        # f_cov = torch.mm(x.transpose(0, 1), x).float()  # d * d
        U, S, V = torch.linalg.svd(0.99 * f_cov + 0.01 * eye)
        diag = torch.diag(1.0 / torch.sqrt(S + 1e-5))
        rotate_mtx = torch.mm(diag, U.transpose(0, 1)).detach()  # d * d
        return torch.mm(rotate_mtx, x.transpose(0, 1)).transpose(0, 1)
//  [python]
def zca_forward(self, x):
        batch_size, feature_dim = x.size()
        eps = 0.1
        f_cov = (torch.mm(x.transpose(0, 1), x) / (batch_size - 1)).float()  # d * d
        eye = torch.eye(feature_dim).float().to(f_cov.device)
        # f_cov = torch.mm(x.transpose(0, 1), x).float()  # d * d
        U, S, V = torch.svd(0.9 * f_cov + 0.1 * eye)
        diag = torch.diag(1.0 / torch.sqrt(S+1e-5))
        rotate_mtx = torch.mm(torch.mm(U, diag), U.transpose(0, 1)).detach() # d * d
        return torch.mm(rotate_mtx, x.transpose(0, 1)).transpose(0, 1)