jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def fit(data, models_posterior):
    # pts: B x C x H x W / B x C x N
    # models_likelihood: B x K x H x W
    # means, stds: B x K x C
    # out: argmax(P(H|E)) B x K x C

    pts = data["pts"]
    B, K, H, W = models_posterior.shape
    # models_posterior = (models_posterior > 0.5) * 1.0
    # B x K x 3
    centroid = (models_posterior[:, :, None] * pts[:, None]).flatten(3).sum(dim=3) / (
        models_posterior[:, :, None].flatten(3).sum(dim=3) + 1e-8
    )
    # B x K x 3 x N
    dev = (
        models_posterior[:, :, None] * (pts[:, None] - centroid[:, :, :, None, None])
    ).flatten(3)
    dev = dev.permute(0, 1, 3, 2)
    # B x K x N x 3

    U, stds, axesT = torch.linalg.svd(dev)
    # B x K x 3
    axes = axesT.permute(0, 1, 3, 2)

    # B x K x N x 3
    dev_proj = dev @ axes
    # B x K x 3
    dev_max, _ = dev_proj.abs().max(dim=2)
    measures = dev_max * 2
    params = {}
    params["center"] = centroid
    params["axes"] = axes
    params["measures"] = measures
    params["stds"] = stds
    params["pts"] = []

    for b in range(B):
        params["pts"].append([])
        for k in range(K):
            params["pts"][b].append(pts[b, :, models_posterior[b, k] > 0.5])

    o4visual.visualize_boxes3d(params)
    # +- 1 std: 68.27%
    # +- 2 std: 95.45%
    # +- 3 std: 99.73%
    return params
//  [python]
def fit_plane_tls_torch(points: np.ndarray, device="cuda", max_n_points=20000):
    """
    ref:
        https://github.com/ajhynes7/scikit-spatial/blob/master/src/skspatial/objects/plane.py
        https://github.com/ajhynes7/scikit-spatial/blob/master/src/skspatial/objects/points.py
    """
    if len(points) > max_n_points:  # subsample to avoid OOM
        points = points[np.random.choice(len(points), max_n_points, replace=False)]

    points = torch.tensor(points, dtype=torch.float32, device=device)
    points_centroid = points.mean(0, keepdim=True)
    points_centered = points - points_centroid
    logger.debug(f"TLS SVD input shape: {points_centered.T.shape} (large inputs might cause OOM)")

    # FIXME: limit maximum point numbers to avoid OOM caused by svd
    U, _, _ = torch.svd(points_centered.T)
    normal = U[:, 2].cpu().numpy()
    centroid = points_centroid[0].cpu().numpy()

    plane = SKPlane(centroid, normal)
    params = np.array(plane.cartesian())
    _scale = np.linalg.norm(params[:3])
    params = params / _scale
    return params
//  [python]
def fit_plane_tls_torch(points: np.ndarray, device="cuda", max_n_points=20000):
    """
    ref:
        https://github.com/ajhynes7/scikit-spatial/blob/master/src/skspatial/objects/plane.py
        https://github.com/ajhynes7/scikit-spatial/blob/master/src/skspatial/objects/points.py
    """
    if len(points) > max_n_points:  # subsample to avoid OOM
        points = points[np.random.choice(len(points), max_n_points, replace=False)]

    points = torch.tensor(points, dtype=torch.float32, device=device)
    points_centroid = points.mean(0, keepdim=True)
    points_centered = points - points_centroid
    logger.debug(f"TLS SVD input shape: {points_centered.T.shape} (large inputs might cause OOM)")

    # FIXME: limit maximum point numbers to avoid OOM caused by svd
    U, _, _ = torch.linalg.svd(points_centered.T)
    normal = U[:, 2].cpu().numpy()
    centroid = points_centroid[0].cpu().numpy()

    plane = SKPlane(centroid, normal)
    params = np.array(plane.cartesian())
    _scale = np.linalg.norm(params[:3])
    params = params / _scale
    return params
//  [python]
def kbash_alignment(self, coor1, coor2):
        A = ( coor2 - coor2.mean(dim=0) ).t() @ ( coor1 - coor1.mean(dim=0) )
        U, S, Vt = torch.svd(A)
        rotation = ( U @ torch.diag( torch.Tensor([1, 1, torch.sign(torch.det(A))]).to(coor1.device) ) ) @ Vt
        translation = coor2.mean(dim=0) - torch.t(rotation @ coor1.mean(dim=0, keepdim=True).t())

        rotation = rotation.detach()
        translation = translation.detach()

        coor1 = coor1 @ rotation.t() + translation
        return coor1
//  [python]
def kbash_alignment(self, coor1, coor2):
        A = ( coor2 - coor2.mean(dim=0) ).t() @ ( coor1 - coor1.mean(dim=0) )
        U, S, Vt = torch.linalg.svd(A)
        rotation = ( U @ torch.diag( torch.Tensor([1, 1, torch.sign(torch.det(A))]).to(coor1.device) ) ) @ Vt
        translation = coor2.mean(dim=0) - torch.t(rotation @ coor1.mean(dim=0, keepdim=True).t())

        rotation = rotation.detach()
        translation = translation.detach()

        coor1 = coor1 @ rotation.t() + translation
        return coor1
//  [python]
def _torch_linalg_svdvals(input: Tensor) -> Tensor:
    """Helper function to make torch.linalg.svdvals 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, Tensor):
        raise AssertionError(f"Input must be Tensor. Got: {type(input)}.")
    dtype: torch.dtype = input.dtype
    if dtype not in (torch.float32, torch.float64):
        dtype = torch.float32

    if TYPE_CHECKING:
        # TODO: remove this branch when kornia relies on torch >= 1.10
        out: Tensor
    elif torch_version_ge(1, 10):
        out = torch.linalg.svdvals(input.to(dtype))
    else:
        # TODO: remove this branch when kornia relies on torch >= 1.10
        _, out, _ = torch.svd(input.to(dtype))
    return out.to(input.dtype)
//  [python]
def _torch_linalg_svdvals(input: Tensor) -> Tensor:
    """Helper function to make torch.linalg.svdvals 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, Tensor):
        raise AssertionError(f"Input must be Tensor. Got: {type(input)}.")
    dtype: torch.dtype = input.dtype
    if dtype not in (torch.float32, torch.float64):
        dtype = torch.float32

    if TYPE_CHECKING:
        # TODO: remove this branch when kornia relies on torch >= 1.10
        out: Tensor
    elif torch_version_ge(1, 10):
        out = torch.linalg.svdvals(input.to(dtype))
    else:
        # TODO: remove this branch when kornia relies on torch >= 1.10
        _, out, _ = torch.linalg.svd(input.to(dtype))
    return out.to(input.dtype)
//  [python]
def _invert_svd(matrix: torch.Tensor, threshold: float = 1e-12) -> torch.Tensor:
    """Invert a matrix using SVD.

    Args:
        matrix: The matrix to invert.
        threshold: The threshold below which singular values are considered zero.

    Returns:
        The inverted matrix.
    """
    u, s, vh = torch.svd(matrix)

    non_zero_idxs = s > threshold

    s_inverse = torch.zeros_like(s)
    s_inverse[non_zero_idxs] = 1.0 / s[non_zero_idxs]

    return vh.T @ torch.diag(s_inverse) @ u.T
//  [python]
def _invert_svd(matrix: torch.Tensor, threshold: float = 1e-12) -> torch.Tensor:
    """Invert a matrix using SVD.

    Args:
        matrix: The matrix to invert.
        threshold: The threshold below which singular values are considered zero.

    Returns:
        The inverted matrix.
    """
    u, s, vh = torch.linalg.svd(matrix)

    non_zero_idxs = s > threshold

    s_inverse = torch.zeros_like(s)
    s_inverse[non_zero_idxs] = 1.0 / s[non_zero_idxs]

    return vh.T @ torch.diag(s_inverse) @ u.T
//  [python]
def align_3d_points_torch(src, dst):
    src_center = src.mean(dim = 0)
    dst_center = dst.mean(dim = 0)
    src_norm = src - src_center
    dst_norm = dst - dst_center
    h_matrix = torch.mm(src_norm.T, dst_norm)

    U, S, Vh = torch.svd(h_matrix)

    R = torch.mm(U, Vh)
    t = dst_center - torch.mm(R.T, src_center.unsqueeze(1)).squeeze()

    return R, t
//  [python]
def align_3d_points_torch(src, dst):
    src_center = src.mean(dim = 0)
    dst_center = dst.mean(dim = 0)
    src_norm = src - src_center
    dst_norm = dst - dst_center
    h_matrix = torch.mm(src_norm.T, dst_norm)

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

    R = torch.mm(U, Vh)
    t = dst_center - torch.mm(R.T, src_center.unsqueeze(1)).squeeze()

    return R, t
//  [python]
def _fit(self, x, y=None):
        try:
            # Future-proofing of SVD, different import point
            # Same result as `torch.linalg.svd`.
            _, _, V = torch.svd(x, **self.init_args)
        except AttributeError:
            # Legacy SVD. Same result as `torch.linalg.svd`.
            _, _, V = torch.svd(x, **self.init_args)

        # Numerical instability: The generated torch tensors are not guaranteed to be unique
        # Specifically, the `V` tensor can be a matrix M, or its negative (-M), depending on the underlying system
        # See: https://pytorch.org/docs/stable/generated/torch.linalg.svd.html#torch.linalg.svd
        # To ensure consistency across systems, we can arbitrarily select the version of the matrix that is positive
        # in the very first index.
        init_value = torch.flatten(V.data).detach().cpu().numpy()[0]
        if init_value < 0:
            V = -V

        if self.n_components:
            V = V[:, :self.n_components]
            if V.shape[1] < self.n_components:
                V = torch.cat((V, torch.zeros(V.shape[0], self.n_components - V.shape[1], requires_grad=True,
                                              device=device)), dim=1)
        self.V = torch.nn.Parameter(V)
//  [python]
def _fit(self, x, y=None):
        try:
            # Future-proofing of SVD, different import point
            # Same result as `torch.svd`.
            _, _, Vh = torch.linalg.svd(x, **self.init_args)
            V = Vh.transpose(-2, -1).conj()
        except AttributeError:
            # Legacy SVD. Same result as `torch.linalg.svd`.
            _, _, V = torch.svd(x, **self.init_args)

        # Numerical instability: The generated torch tensors are not guaranteed to be unique
        # Specifically, the `V` tensor can be a matrix M, or its negative (-M), depending on the underlying system
        # See: https://pytorch.org/docs/stable/generated/torch.linalg.svd.html#torch.linalg.svd
        # To ensure consistency across systems, we can arbitrarily select the version of the matrix that is positive
        # in the very first index.
        init_value = torch.flatten(V.data).detach().cpu().numpy()[0]
        if init_value < 0:
            V = -V

        if self.n_components:
            V = V[:, :self.n_components]
            if V.shape[1] < self.n_components:
                V = torch.cat((V, torch.zeros(V.shape[0], self.n_components - V.shape[1], requires_grad=True,
                                              device=device)), dim=1)
        self.V = torch.nn.Parameter(V)
//  [python]
def _fit(self, X):
        if self.n_components is None:
            n_components = min(X.shape)
        else:
            n_components = self.n_components

        n_samples, n_features = X.shape
        if n_components > min(X.shape):
            raise ValueError(
                f"n_components should be <= min(n_samples: {n_samples}, n_features: {n_features})"
            )

        self.mean_ = torch.mean(X, axis=0)
        X -= self.mean_

        if self.fit == "full":
            U, S, Vt = torch.svd(X, some=False)
            # flip eigenvectors' sign to enforce deterministic output
            U, Vt = svd_flip(U, Vt)
        elif self.fit == "lowrank":
            U, S, Vt = torch.pca_lowrank(X)

        self.components_ = Vt[:n_components]
        self.n_components_ = n_components

        return U, S, Vt
//  [python]
def _fit(self, X):
        if self.n_components is None:
            n_components = min(X.shape)
        else:
            n_components = self.n_components

        n_samples, n_features = X.shape
        if n_components > min(X.shape):
            raise ValueError(
                f"n_components should be <= min(n_samples: {n_samples}, n_features: {n_features})"
            )

        self.mean_ = torch.mean(X, axis=0)
        X -= self.mean_

        if self.fit == "full":
            U, S, Vt = torch.linalg.svd(X, full_matrices=False)
            # flip eigenvectors' sign to enforce deterministic output
            U, Vt = svd_flip(U, Vt)
        elif self.fit == "lowrank":
            U, S, Vt = torch.pca_lowrank(X)

        self.components_ = Vt[:n_components]
        self.n_components_ = n_components

        return U, S, Vt
//  [python]
def convert_to_standard_lora(self):
        # convert down and up matrix to linear layer, nn.Linear call F.linear
        with torch.no_grad():
            # merge two lora weight
            new_weight = torch.mm(self.up_frozen.weight, self.down_frozen.weight) #+ \
                         # torch.mm(self.up_trainable.weight, self.down_trainable.weight)

            # torch.svd
            U, s, Vt = torch.svd(new_weight)
            print("shape====================>")
            print(new_weight.shape)
            print(U.shape)
            print(Vt.shape)

            # Select the first r singular values
            r = self.rank_f + self.rank_t
            Ur = U[:, :r]
            sr = s[:r]
            Vtr = Vt[:r, :]

            # set the weights of down and up matrix
            up = Ur * sr
            down = Vtr

            if not up.is_contiguous():
                up = up.contiguous()
            if not down.is_contiguous():
                down = down.contiguous()
            down = nn.Parameter(down)
            up = nn.Parameter(up)

            # convert down and up matrix to linear layer
            self.down = nn.Linear(self.in_features, r, bias=False)
            self.down.weight = down
            self.up = nn.Linear(r, self.out_features, bias=False)
            self.up.weight = nn.Parameter(up)

        # Remove redundant layers and variables
        del self.down_frozen
        del self.up_frozen
        del self.down_trainable
        del self.up_trainable
//  [python]
def convert_to_standard_lora(self):
        # convert down and up matrix to linear layer, nn.Linear call F.linear
        with torch.no_grad():
            # merge two lora weight
            new_weight = torch.mm(self.up_frozen.weight, self.down_frozen.weight) #+ \
                         # torch.mm(self.up_trainable.weight, self.down_trainable.weight)

            # torch.svd
            U, s, Vt = torch.linalg.svd(new_weight, full_matrices=False)
            print("shape====================>")
            print(new_weight.shape)
            print(U.shape)
            print(Vt.shape)

            # Select the first r singular values
            r = self.rank_f + self.rank_t
            Ur = U[:, :r]
            sr = s[:r]
            Vtr = Vt[:r, :]

            # set the weights of down and up matrix
            up = Ur * sr
            down = Vtr

            if not up.is_contiguous():
                up = up.contiguous()
            if not down.is_contiguous():
                down = down.contiguous()
            down = nn.Parameter(down)
            up = nn.Parameter(up)

            # convert down and up matrix to linear layer
            self.down = nn.Linear(self.in_features, r, bias=False)
            self.down.weight = down
            self.up = nn.Linear(r, self.out_features, bias=False)
            self.up.weight = nn.Parameter(up)

        # Remove redundant layers and variables
        del self.down_frozen
        del self.up_frozen
        del self.down_trainable
        del self.up_trainable
//  [python]
def __init__(self, in_features: int, out_features: int, name: str, config: dict, **kwargs):
        self.r = {}
        self.lora_alpha = {}
        self.scaling = {}
        self.lora_dropout = nn.ModuleDict({})
        self.lora_A = nn.ModuleDict({})
        self.lora_B = nn.ModuleDict({})
        self.lora_E = nn.ModuleDict({}) 
        # For Embedding layer
        self.lora_embedding_A = nn.ParameterDict({})
        self.lora_embedding_B = nn.ParameterDict({})
        self.lora_embedding_E = nn.ParameterDict({})
        # Mark the weight as unmerged
        self.merged = False
        self.disable_adapters = False
        self.in_features = in_features
        self.out_features = out_features
        self.kwargs = kwargs
        self.config = config
        
        weight_copy = self.weight.clone()
        
        if ("V" in self.config.update_rule["A_init"]) or ("U" in self.config.update_rule["B_init"]):
            U, S, V = torch.svd(weight_copy)
        else:
            U, S, V = None, None, None    
        self.U_init = U
        self.S_init = S
        self.V_init = V        
        
        if (self.config.update_rule["A_init"] == "rand") or (self.config.update_rule["B_init"] == "rand"):
            n_w, m_w = weight_copy.size()
            random_w = torch.rand(n_w, m_w) - 0.5
            
            self.U_rand, self.S_rand, self.V_rand = torch.svd(random_w)
        else:
            self.U_rand, self.S_rand, self.V_rand = None, None, None
//  [python]
def __init__(self, in_features: int, out_features: int, name: str, config: dict, **kwargs):
        self.r = {}
        self.lora_alpha = {}
        self.scaling = {}
        self.lora_dropout = nn.ModuleDict({})
        self.lora_A = nn.ModuleDict({})
        self.lora_B = nn.ModuleDict({})
        self.lora_E = nn.ModuleDict({}) 
        # For Embedding layer
        self.lora_embedding_A = nn.ParameterDict({})
        self.lora_embedding_B = nn.ParameterDict({})
        self.lora_embedding_E = nn.ParameterDict({})
        # Mark the weight as unmerged
        self.merged = False
        self.disable_adapters = False
        self.in_features = in_features
        self.out_features = out_features
        self.kwargs = kwargs
        self.config = config
        
        weight_copy = self.weight.clone()
        
        if ("V" in self.config.update_rule["A_init"]) or ("U" in self.config.update_rule["B_init"]):
            U, S, V = torch.linalg.svd(weight_copy)
        else:
            U, S, V = None, None, None    
        self.U_init = U
        self.S_init = S
        self.V_init = V        
        
        if (self.config.update_rule["A_init"] == "rand") or (self.config.update_rule["B_init"] == "rand"):
            n_w, m_w = weight_copy.size()
            random_w = torch.rand(n_w, m_w) - 0.5
            
            self.U_rand, self.S_rand, self.V_rand = torch.linalg.svd(random_w)
        else:
            self.U_rand, self.S_rand, self.V_rand = None, None, None
//  [python]
def canonical_stiefel_case():
    torch.manual_seed(42)
    shape = manifold_shapes[geoopt.manifolds.CanonicalStiefel]
    ex = torch.randn(*shape)
    ev = torch.randn(*shape)
    u, _, v = torch.svd(ex, some=False)
    x = torch.einsum("...ik,...kj->...ij", u, v)
    v = ev - x @ ev.t() @ x
    manifold = geoopt.manifolds.CanonicalStiefel()
    x = geoopt.ManifoldTensor(x, manifold=manifold)
    case = UnaryCase(shape, x, ex, v, ev, manifold)
    yield case
//  [python]
def canonical_stiefel_case():
    torch.manual_seed(42)
    shape = manifold_shapes[geoopt.manifolds.CanonicalStiefel]
    ex = torch.randn(*shape)
    ev = torch.randn(*shape)
    u, _, v = torch.linalg.svd(ex, full_matrices=False)
    x = torch.einsum("...ik,...kj->...ij", u, v)
    v = ev - x @ ev.t() @ x
    manifold = geoopt.manifolds.CanonicalStiefel()
    x = geoopt.ManifoldTensor(x, manifold=manifold)
    case = UnaryCase(shape, x, ex, v, ev, manifold)
    yield case
//  [python]
def punish_wight(wo_batch, latent_size, alpha, method):
    if method == 'weight':
        wo_batch *= alpha
    elif method in ['alpha', 'beta', 'delete', 'soft-weight']:
        u, s, vh = torch.svd(wo_batch)
        u = u[:,:latent_size]
        zero_idx = int(latent_size * alpha)
        if method == 'alpha':
            s[:zero_idx] = 0
        elif method == 'beta':
            s[zero_idx:] = 0
        elif method == 'delete':
            s = s[zero_idx:] if zero_idx < latent_size else torch.zeros(latent_size).to(s.device)
            u = u[:, zero_idx:] if zero_idx < latent_size else u
            vh = vh[zero_idx:, :] if zero_idx < latent_size else vh
        elif method == 'soft-weight':
            if CALC_SIMILARITY:
                _s = s.clone()
                _s[zero_idx:] = 0
                _wo_batch = u @ torch.diag(_s) @ vh
                dist = cdist(wo_batch[:,0].unsqueeze(0).cpu(), _wo_batch[:,0].unsqueeze(0).cpu(), metric='cosine')
                print(f'The distance between the word embedding before and after the punishment: {dist}')
            if alpha == -.001:
                s *= (torch.exp(-.001 * s) * 1.2)  # strengthen objects (our Appendix.F)
            else:
                s *= torch.exp(-alpha*s)  # suppression EOT (our main paper)

        wo_batch = u @ torch.diag(s) @ vh
    else:
        raise ValueError('Unsupported method')
    return wo_batch
//  [python]
def punish_wight(wo_batch, latent_size, alpha, method):
    if method == 'weight':
        wo_batch *= alpha
    elif method in ['alpha', 'beta', 'delete', 'soft-weight']:
        u, s, vh = torch.linalg.svd(wo_batch)
        u = u[:,:latent_size]
        zero_idx = int(latent_size * alpha)
        if method == 'alpha':
            s[:zero_idx] = 0
        elif method == 'beta':
            s[zero_idx:] = 0
        elif method == 'delete':
            s = s[zero_idx:] if zero_idx < latent_size else torch.zeros(latent_size).to(s.device)
            u = u[:, zero_idx:] if zero_idx < latent_size else u
            vh = vh[zero_idx:, :] if zero_idx < latent_size else vh
        elif method == 'soft-weight':
            if CALC_SIMILARITY:
                _s = s.clone()
                _s[zero_idx:] = 0
                _wo_batch = u @ torch.diag(_s) @ vh
                dist = cdist(wo_batch[:,0].unsqueeze(0).cpu(), _wo_batch[:,0].unsqueeze(0).cpu(), metric='cosine')
                print(f'The distance between the word embedding before and after the punishment: {dist}')
            if alpha == -.001:
                s *= (torch.exp(-.001 * s) * 1.2)  # strengthen objects (our Appendix.F)
            else:
                s *= torch.exp(-alpha*s)  # suppression EOT (our main paper)

        wo_batch = u @ torch.diag(s) @ vh
    else:
        raise ValueError('Unsupported method')
    return wo_batch
//  [python]
def euclidean_stiefel_case():
    torch.manual_seed(42)
    shape = manifold_shapes[geoopt.manifolds.EuclideanStiefel]
    ex = torch.randn(*shape, dtype=torch.float64)
    ev = torch.randn(*shape, dtype=torch.float64)
    u, s, v = torch.svd(ex)
    x = torch.einsum("...ik,...kj->...ij", u, v)
    nonsym = x.t() @ ev
    v = ev - x @ (nonsym + nonsym.t()) / 2

    manifold = geoopt.manifolds.EuclideanStiefel()
    x = geoopt.ManifoldTensor(x, manifold=manifold)
    case = UnaryCase(shape, x, ex, v, ev, manifold)
    yield case
    manifold = geoopt.manifolds.EuclideanStiefelExact()
    x = geoopt.ManifoldTensor(x, manifold=manifold)
    case = UnaryCase(shape, x, ex, v, ev, manifold)
    yield case
//  [python]
def euclidean_stiefel_case():
    torch.manual_seed(42)
    shape = manifold_shapes[geoopt.manifolds.EuclideanStiefel]
    ex = torch.randn(*shape, dtype=torch.float64)
    ev = torch.randn(*shape, dtype=torch.float64)
    u, _, v = torch.linalg.svd(ex, full_matrices=False)
    x = torch.einsum("...ik,...kj->...ij", u, v)
    nonsym = x.t() @ ev
    v = ev - x @ (nonsym + nonsym.t()) / 2

    manifold = geoopt.manifolds.EuclideanStiefel()
    x = geoopt.ManifoldTensor(x, manifold=manifold)
    case = UnaryCase(shape, x, ex, v, ev, manifold)
    yield case
    manifold = geoopt.manifolds.EuclideanStiefelExact()
    x = geoopt.ManifoldTensor(x, manifold=manifold)
    case = UnaryCase(shape, x, ex, v, ev, manifold)
    yield case
//  [python]
def get_transformation(hm, region_score, batch, v0, v1):
    with torch.no_grad():
        hd0 = scatter_sum(hm[:, :, None] * v0[:, None], batch, dim=0)
        hd1 = scatter_sum(hm[:, :, None] * v1[:, None], batch, dim=0)  # (B, 40, 3)

        B, K, _ = hd0.shape
        seg = torch.max(region_score, 1)[1]  # (B*V,)
        R_batch = []
        for i in range(B):
            center0 = hd0[i]  # (40, 3)
            center1 = hd1[i]  # (40, 3)
            seg_single = seg[batch==i]  # (V,)
            v0_single = v0[batch==i]  # (V, 3)
            v1_single = v1[batch==i]  # (V, 3)
            seg_max = torch.max(seg_single) + 1

            scatter_center0 = center0[seg_single]
            scatter_center1 = center1[seg_single]

            v0_single = v0_single - scatter_center0
            v1_single = v1_single - scatter_center1

            M = torch.bmm(v0_single[:, :, None], v1_single[:, None])  # (V, 3, 3)
            M_sum = scatter_sum(M, seg_single, dim=0)  # (40, 3, 3)
            U, S, Vh = torch.svd(M_sum)
            R_temp = torch.bmm(U, Vh).permute(0, 2, 1)  # (40, 3, 3)
            R = torch.eye(3)[None].repeat(K, 1, 1).float().to(R_temp.device)
            R[:seg_max] = R_temp
            R = rotation_matrix_to_quaternion(R, order=wxyz)
            R_batch.append(R)

        R_batch = torch.stack(R_batch, 0)  # (B, 40, 4)
        t_batch = hd1 - hd0
    return torch.cat((t_batch, R_batch), 2)
//  [python]
def get_transformation(hm, region_score, batch, v0, v1):
    with torch.no_grad():
        hd0 = scatter_sum(hm[:, :, None] * v0[:, None], batch, dim=0)
        hd1 = scatter_sum(hm[:, :, None] * v1[:, None], batch, dim=0)  # (B, 40, 3)

        B, K, _ = hd0.shape
        seg = torch.max(region_score, 1)[1]  # (B*V,)
        R_batch = []
        for i in range(B):
            center0 = hd0[i]  # (40, 3)
            center1 = hd1[i]  # (40, 3)
            seg_single = seg[batch==i]  # (V,)
            v0_single = v0[batch==i]  # (V, 3)
            v1_single = v1[batch==i]  # (V, 3)
            seg_max = torch.max(seg_single) + 1

            scatter_center0 = center0[seg_single]
            scatter_center1 = center1[seg_single]

            v0_single = v0_single - scatter_center0
            v1_single = v1_single - scatter_center1

            M = torch.bmm(v0_single[:, :, None], v1_single[:, None])  # (V, 3, 3)
            M_sum = scatter_sum(M, seg_single, dim=0)  # (40, 3, 3)
            U, S, Vh = torch.linalg.svd(M_sum)
            R_temp = torch.bmm(U, Vh).permute(0, 2, 1)  # (40, 3, 3)
            R = torch.eye(3)[None].repeat(K, 1, 1).float().to(R_temp.device)
            R[:seg_max] = R_temp
            R = rotation_matrix_to_quaternion(R, order=wxyz)
            R_batch.append(R)

        R_batch = torch.stack(R_batch, 0)  # (B, 40, 4)
        t_batch = hd1 - hd0
    return torch.cat((t_batch, R_batch), 2)
//  [python]
def compress_delta(delta_dict, unet, compression_ratio = 0.6):
    compressed_st = {'unet': {}}
    pretrained_st = unet.state_dict()
    st = delta_dict
    if 'modifier_token' in st:
        compressed_st['modifier_token'] = st['modifier_token']
    if 'text_encoder' in st:
        compressed_st['text_encoder'] = st['text_encoder']
    st = st['unet']
    layers = list(st.keys())
    for name in layers:
        if 'to_k' in name or 'to_v' in name:
            W = st[name].detach().to(unet.device)
            Wpretrain = pretrained_st[name]
            deltaW = W - Wpretrain
            u, s, vt = torch.svd(deltaW.cpu()) # OOM if on cuda during training
            explain = 0
            all_ = (s).sum()
            for i, t in enumerate(s):
                explain += t/(all_)
                if explain > compression_ratio:
                    break
            compressed_st['unet'][f'{name}'] = {}
            compressed_st['unet'][f'{name}']['u'] = (u[:, :i]@torch.diag(s)[:i, :i]).clone().cuda()
            compressed_st['unet'][f'{name}']['v'] = vt[:i].clone().cuda()
        else:
            compressed_st['unet'][f'{name}'] = st[name]
    return compressed_st
//  [python]
def compress_delta(delta_dict, unet, compression_ratio = 0.6):
    compressed_st = {'unet': {}}
    pretrained_st = unet.state_dict()
    st = delta_dict
    if 'modifier_token' in st:
        compressed_st['modifier_token'] = st['modifier_token']
    if 'text_encoder' in st:
        compressed_st['text_encoder'] = st['text_encoder']
    st = st['unet']
    layers = list(st.keys())
    for name in layers:
        if 'to_k' in name or 'to_v' in name:
            W = st[name].detach().to(unet.device)
            Wpretrain = pretrained_st[name]
            deltaW = W - Wpretrain
            u, s, vt = torch.linalg.svd(deltaW.cpu()) # OOM if on cuda during training
            explain = 0
            all_ = (s).sum()
            for i, t in enumerate(s):
                explain += t/(all_)
                if explain > compression_ratio:
                    break
            compressed_st['unet'][f'{name}'] = {}
            compressed_st['unet'][f'{name}']['u'] = (u[:, :i]@torch.diag(s)[:i, :i]).clone().cuda()
            compressed_st['unet'][f'{name}']['v'] = vt[:i].clone().cuda()
        else:
            compressed_st['unet'][f'{name}'] = st[name]
    return compressed_st
//  [python]
def _svd(input: torch.Tensor
         ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    # torch.svd style
    U, S, Vh = torch.svd(input, some=False)
    V = Vh.transpose(-2, -1)
    return U, S, V
//  [python]
def _svd(input: torch.Tensor
         ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    # torch.svd style
    U, S, Vh = torch.linalg.svd(input, full_matrices=False)
    V = Vh.transpose(-2, -1)
    return U, S, V
//  [python]
def _estimate_rhs(self, old_pnts, new_pnts):
        # get the neighborhoods
        O = -( old_pnts[self.Lnbrs, :] - old_pnts[:,None,:])*self.Lwgts[:,:,None]
        N = -( new_pnts[self.Lnbrs, :] - new_pnts[:,None,:])*self.Lwgts[:,:,None]
        
        # compute the rotation between the neighborhoods
        cov = N.permute(0,2,1) @ O
        U, s, Vt = torch.svd(cov)

        # fix up the jacobian, if needed
        ones = torch.ones_like(s[:, 0], device=self.device)
        sign = torch.sign(torch.det(U @ Vt))
        s_new = torch.stack((ones,ones,sign), dim=-1)
        R = U @ (s_new[...,None]*Vt)

        # use average rotation for the edge to rotate the edge vectors
        d_new = torch.sum( ((R[self.Lnbrs, ...] + R[:,None,:,:]) @ O[..., None])[...,0]*(0.5*self.Lwgts[:,:,None]), dim=1 )
        return self._mul_L(d_new)
//  [python]
def _estimate_rhs(self, old_pnts, new_pnts):
        # get the neighborhoods
        O = -( old_pnts[self.Lnbrs, :] - old_pnts[:,None,:])*self.Lwgts[:,:,None]
        N = -( new_pnts[self.Lnbrs, :] - new_pnts[:,None,:])*self.Lwgts[:,:,None]
        
        # compute the rotation between the neighborhoods
        cov = N.permute(0,2,1) @ O
        U, s, Vt = torch.linalg.svd(cov)

        # fix up the jacobian, if needed
        ones = torch.ones_like(s[:, 0], device=self.device)
        sign = torch.sign(torch.det(U @ Vt))
        s_new = torch.stack((ones,ones,sign), dim=-1)
        R = U @ (s_new[...,None]*Vt)

        # use average rotation for the edge to rotate the edge vectors
        d_new = torch.sum( ((R[self.Lnbrs, ...] + R[:,None,:,:]) @ O[..., None])[...,0]*(0.5*self.Lwgts[:,:,None]), dim=1 )
        return self._mul_L(d_new)
//  [python]
def pinv_naive(A):
    device = A.device
    dtype = A.dtype
    U, S, Vh = torch.svd(A, some=False)
    threshold = torch.max(S) * 1e-5
    S_mask = S[S > threshold]
    S_pinv = torch.cat([1.0 / S_mask, torch.full([S.numel() - S_mask.numel()], 0.0, device=device, dtype=dtype)], 0)
    A_pinv = Vh.t() @ S_pinv.diag() @ U.t()
    return A_pinv
//  [python]
def pinv_naive(A):
    device = A.device
    dtype = A.dtype
    # U, S, V = torch.svd(A, some=False)
    U, S, Vh = torch.linalg.svd(A, full_matrices=True)
    threshold = torch.max(S) * 1e-5
    # S_pinv = torch.where(S > threshold, 1/S, torch.zeros_like(S))
    S_mask = S[S > threshold]
    S_pinv = torch.cat([1.0 / S_mask, torch.full([S.numel() - S_mask.numel()], 0.0, device=device, dtype=dtype)], 0)
    # A_pinv = V @ S_pinv.diag() @ U.t()
    A_pinv = Vh.t() @ S_pinv.diag() @ U.t()
    return A_pinv
//  [python]
def SVD(self) -> 'Pose':
        U, _, Vt = torch.svd(self.R)
        R = U @ Vt
        return self.__class__.from_Rt(R, self.t)
//  [python]
def SVD(self) -> 'Pose':
        U, _, Vt = torch.linalg.svd(self.R)
        R = U @ Vt
        return self.__class__.from_Rt(R, self.t)
//  [python]
def svd(tensor, like=None, **kwargs):
    """Compute the singular value decomposition of a tensor in each interface.

    The singular value decomposition for a matrix :math:`A` consist of three matrices :math:`S`,
    :math:`U` and :math:`V_h`, such that:

    .. math::

        A = U . Diag(S) . V_h

    Args:
        tensor (tensor_like): input tensor
        compute_uv (bool):  if ``True``, the full decomposition is returned


    Returns:
        :math:`S`, :math:`U` and :math:`V_h` or :math:`S`: full decomposition
        if ``compute_uv`` is ``True`` or ``None``, or only the singular values
        if ``compute_uv`` is ``False``
    """
    if like == "tensorflow":
        from tensorflow.linalg import svd, adjoint

        # Tensorflow results need some post-processing to keep it similar to other frameworks.

        if kwargs.get("compute_uv", True):
            S, U, V = svd(tensor, **kwargs)
            return U, S, adjoint(V)
        return svd(tensor, **kwargs)

    if like == "jax":
        from jax.numpy.linalg import svd

    elif like == "torch":
        if kwargs.get("compute_uv", True) is False:
            from torch import svdvals as svd
        else:
            from torch import svd
        if kwargs.get("compute_uv", None) is not None:
            kwargs.pop("compute_uv")

    else:
        from numpy.linalg import svd

    return svd(tensor, **kwargs)
//  [python]
def svd(tensor, like=None, **kwargs):
    """Compute the singular value decomposition of a tensor in each interface.

    The singular value decomposition for a matrix :math:`A` consist of three matrices :math:`S`,
    :math:`U` and :math:`V_h`, such that:

    .. math::

        A = U . Diag(S) . V_h

    Args:
        tensor (tensor_like): input tensor
        compute_uv (bool):  if ``True``, the full decomposition is returned


    Returns:
        :math:`S`, :math:`U` and :math:`V_h` or :math:`S`: full decomposition
        if ``compute_uv`` is ``True`` or ``None``, or only the singular values
        if ``compute_uv`` is ``False``
    """
    if like == "tensorflow":
        from tensorflow.linalg import svd, adjoint

        # Tensorflow results need some post-processing to keep it similar to other frameworks.

        if kwargs.get("compute_uv", True):
            S, U, V = svd(tensor, **kwargs)
            return U, S, adjoint(V)
        return svd(tensor, **kwargs)

    if like == "jax":
        from jax.numpy.linalg import svd

    elif like == "torch":
        # Torch is deprecating torch.svd() in favour of torch.linalg.svd().
        # The new UI is slightly different and breaks the logic for the multi dispatching.
        # This small workaround restores the compute_uv control argument.
        if kwargs.get("compute_uv", True) is False:
            from torch.linalg import svdvals as svd
        else:
            from torch.linalg import svd
        if kwargs.get("compute_uv", None) is not None:
            kwargs.pop("compute_uv")

    else:
        from numpy.linalg import svd

    return svd(tensor, **kwargs)
//  [python]
def transformation_loss(transformation, hm, region_score, batch, v0, v1, criterion):
    hd0 = scatter_sum(hm[:, :, None] * v0[:, None], batch, dim=0)
    hd1 = scatter_sum(hm[:, :, None] * v1[:, None], batch, dim=0)  # (B, 40, 3)
    r_weight = 0.5
    with torch.no_grad():
        B, K, _ = hd0.shape
        seg = torch.max(region_score, 1)[1]  # (B*V,)
        R_batch = []
        for i in range(B):
            center0 = hd0[i]  # (40, 3)
            center1 = hd1[i]  # (40, 3)
            seg_single = seg[batch==i]  # (V,)
            v0_single = v0[batch==i]  # (V, 3)
            v1_single = v1[batch==i]  # (V, 3)
            seg_max = torch.max(seg_single) + 1

            scatter_center0 = center0[seg_single]
            scatter_center1 = center1[seg_single]

            v0_single = v0_single - scatter_center0
            v1_single = v1_single - scatter_center1

            M = torch.bmm(v0_single[:, :, None], v1_single[:, None])  # (V, 3, 3)
            M_sum = scatter_sum(M, seg_single, dim=0)  # (40, 3, 3)
            try:
                U, S, Vh = torch.svd(M_sum)
                # u, s, vh = np.linalg.svd(M_sum.cpu().numpy())
                # U = torch.from_numpy(u).float().to(M_sum.device)
                # S = torch.from_numpy(s).float().to(M_sum.device)
                # Vh = torch.from_numpy(vh).float().to(M_sum.device)
            except RuntimeError as e:
                print(e)
                r_weight = 0
                R_batch.append(torch.eye(3)[None].repeat(K, 1, 1).float().to(M_sum.device))
                continue
            R_temp = torch.bmm(U, Vh).permute(0, 2, 1)  # (40, 3, 3)
            R = torch.eye(3)[None].repeat(K, 1, 1).float().to(R_temp.device)
            R[:seg_max] = R_temp
            R = rotation_matrix_to_quaternion(R, order=wxyz)
            R_batch.append(R)

        R_batch = torch.stack(R_batch, 0)  # (B, 40, 3, 3)

        # R_batch = R_batch.detach().cpu().numpy()
        # Rs = R_batch.reshape([-1, 3, 3])
        # R_batch = []
        # for R in Rs:
        #     R_batch.append(cv2.Rodrigues(R)[0])
        # R_batch = np.stack(R_batch, 0)
        # R_batch = R_batch.reshape([B, K, 3])
        # R_batch = torch.from_numpy(R_batch).float().to(hd0.device)

        t_batch = hd1 - hd0
    pred_t = transformation[:, :, :3]
    pred_R = transformation[:, :, 3:]
    if r_weight == 0:
        return criterion(t_batch, pred_t)
    else:
        return quaternion_loss(R_batch, pred_R) * r_weight + criterion(t_batch, pred_t)
//  [python]
def transformation_loss(transformation, hm, region_score, batch, v0, v1, criterion):
    hd0 = scatter_sum(hm[:, :, None] * v0[:, None], batch, dim=0)
    hd1 = scatter_sum(hm[:, :, None] * v1[:, None], batch, dim=0)  # (B, 40, 3)
    r_weight = 0.5
    with torch.no_grad():
        B, K, _ = hd0.shape
        seg = torch.max(region_score, 1)[1]  # (B*V,)
        R_batch = []
        for i in range(B):
            center0 = hd0[i]  # (40, 3)
            center1 = hd1[i]  # (40, 3)
            seg_single = seg[batch==i]  # (V,)
            v0_single = v0[batch==i]  # (V, 3)
            v1_single = v1[batch==i]  # (V, 3)
            seg_max = torch.max(seg_single) + 1

            scatter_center0 = center0[seg_single]
            scatter_center1 = center1[seg_single]

            v0_single = v0_single - scatter_center0
            v1_single = v1_single - scatter_center1

            M = torch.bmm(v0_single[:, :, None], v1_single[:, None])  # (V, 3, 3)
            M_sum = scatter_sum(M, seg_single, dim=0)  # (40, 3, 3)
            try:
                U, S, Vh = torch.linalg.svd(M_sum)
                # u, s, vh = np.linalg.svd(M_sum.cpu().numpy())
                # U = torch.from_numpy(u).float().to(M_sum.device)
                # S = torch.from_numpy(s).float().to(M_sum.device)
                # Vh = torch.from_numpy(vh).float().to(M_sum.device)
            except RuntimeError as e:
                print(e)
                r_weight = 0
                R_batch.append(torch.eye(3)[None].repeat(K, 1, 1).float().to(M_sum.device))
                continue
            R_temp = torch.bmm(U, Vh).permute(0, 2, 1)  # (40, 3, 3)
            R = torch.eye(3)[None].repeat(K, 1, 1).float().to(R_temp.device)
            R[:seg_max] = R_temp
            R = rotation_matrix_to_quaternion(R, order=wxyz)
            R_batch.append(R)

        R_batch = torch.stack(R_batch, 0)  # (B, 40, 3, 3)

        # R_batch = R_batch.detach().cpu().numpy()
        # Rs = R_batch.reshape([-1, 3, 3])
        # R_batch = []
        # for R in Rs:
        #     R_batch.append(cv2.Rodrigues(R)[0])
        # R_batch = np.stack(R_batch, 0)
        # R_batch = R_batch.reshape([B, K, 3])
        # R_batch = torch.from_numpy(R_batch).float().to(hd0.device)

        t_batch = hd1 - hd0
    pred_t = transformation[:, :, :3]
    pred_R = transformation[:, :, 3:]
    if r_weight == 0:
        return criterion(t_batch, pred_t)
    else:
        return quaternion_loss(R_batch, pred_R) * r_weight + criterion(t_batch, pred_t)
//  [python]
def get_orthogonal_matrix(self, weights, rank, type):
        module_params = weights

        if module_params.data.dtype != torch.float:
            float_data = False
            original_type = module_params.data.dtype
            original_device = module_params.data.device
            matrix = module_params.data.float()
        else:
            float_data = True
            matrix = module_params.data
            
        U, s, Vh = torch.svd(matrix)
        
        #make the smaller matrix always to be orthogonal matrix
        if type=='right':
            A = U[:, :rank] @ torch.diag(s[:rank])
            B = Vh[:rank, :]
            
            if not float_data:
                B = B.to(original_device).type(original_type)
            return B
        elif type=='left':
            A = U[:, :rank]
            B = torch.diag(s[:rank]) @ Vh[:rank, :]
            if not float_data:
                A = A.to(original_device).type(original_type)
            return A
        elif type=='full':
            A = U[:, :rank]
            B = Vh[:rank, :]
            if not float_data:
                A = A.to(original_device).type(original_type)
                B = B.to(original_device).type(original_type)
            return [A, B]
        else:
            raise ValueError('type should be left, right or full')
//  [python]
def get_orthogonal_matrix(self, weights, rank, type):
        module_params = weights

        if module_params.data.dtype != torch.float:
            float_data = False
            original_type = module_params.data.dtype
            original_device = module_params.data.device
            matrix = module_params.data.float()
        else:
            float_data = True
            matrix = module_params.data
            
        U, s, Vh = torch.linalg.svd(matrix, full_matrices = False)
        
        #make the smaller matrix always to be orthogonal matrix
        if type=='right':
            A = U[:, :rank] @ torch.diag(s[:rank])
            B = Vh[:rank, :]
            
            if not float_data:
                B = B.to(original_device).type(original_type)
            return B
        elif type=='left':
            A = U[:, :rank]
            B = torch.diag(s[:rank]) @ Vh[:rank, :]
            if not float_data:
                A = A.to(original_device).type(original_type)
            return A
        elif type=='full':
            A = U[:, :rank]
            B = Vh[:rank, :]
            if not float_data:
                A = A.to(original_device).type(original_type)
                B = B.to(original_device).type(original_type)
            return [A, B]
        else:
            raise ValueError('type should be left, right or full')
//  [python]
def get_svd_u_vec(model, toxic_vector, topk_sorted_score, U_idx):
    """
    Get the svd U vector
    toxic_vector: toxic_vector [d_model]
    topk_sorted_score: (int) vectors we want to get
    U_idx: which u vec
    """
    scores = []
    for layer in range(model.config.n_layer):
        # mlp_outs = model.blocks[layer].mlp.W_out
        # [d_mlp, d_model]
        mlp_outs = model.transformer.h[layer].mlp.c_proj.weight
        cos_sims = F.cosine_similarity(
            mlp_outs, toxic_vector.unsqueeze(0), dim=1
        )
        _topk = cos_sims.topk(k=100)
        _values = [x.item() for x in _topk.values]
        _idxs = [x.item() for x in _topk.indices]
        topk = list(zip(_values, _idxs, [layer] * _topk.indices.shape[0]))
        scores.extend(topk)

    sorted_scores = sorted(scores, key=lambda x: x[0], reverse=True)
    top_vecs = [
        # model.blocks[x[2]].mlp.W_out[x[1]]
        model.transformer.h[x[2]].mlp.c_proj.weight[x[1]]
        for x in sorted_scores[:topk_sorted_score]
    ]
    top_vecs = [x / x.norm() for x in top_vecs]
    _top_vecs = torch.stack(top_vecs)

    svd = torch.svd(_top_vecs.transpose(0, 1))
    svd_U = svd[0].transpose(0, 1)
    return svd_U[U_idx]
//  [python]
def get_svd_u_vec(model, toxic_vector, topk_sorted_score, U_idx):
    """
    Get the svd U vector
    toxic_vector: toxic_vector [d_model]
    topk_sorted_score: (int) vectors we want to get
    U_idx: which u vec
    """
    scores = []
    for layer in range(model.config.n_layer):
        # mlp_outs = model.blocks[layer].mlp.W_out
        # [d_mlp, d_model]
        mlp_outs = model.transformer.h[layer].mlp.c_proj.weight
        cos_sims = F.cosine_similarity(
            mlp_outs, toxic_vector.unsqueeze(0), dim=1
        )
        _topk = cos_sims.topk(k=100)
        _values = [x.item() for x in _topk.values]
        _idxs = [x.item() for x in _topk.indices]
        topk = list(zip(_values, _idxs, [layer] * _topk.indices.shape[0]))
        scores.extend(topk)

    sorted_scores = sorted(scores, key=lambda x: x[0], reverse=True)
    top_vecs = [
        # model.blocks[x[2]].mlp.W_out[x[1]]
        model.transformer.h[x[2]].mlp.c_proj.weight[x[1]]
        for x in sorted_scores[:topk_sorted_score]
    ]
    top_vecs = [x / x.norm() for x in top_vecs]
    _top_vecs = torch.stack(top_vecs)

    svd = torch.linalg.svd(_top_vecs.transpose(0, 1))
    svd_U = svd.U.transpose(0, 1)
    return svd_U[U_idx]
//  [python]
def coral(source_features, val_target_features, unlabeled_target_features, labeled_target_features, k, save_path):
    x_s = torch.cat((source_features.to(device), labeled_target_features.to(device)))
    x_t = unlabeled_target_features.to(device)
    x_tv = val_target_features.to(device)
    
    # PCA
    x = torch.cat((x_s, x_t))
    x = x - x.mean(0)
    
    # below line takes forever, save it
    if os.path.isfile(save_path):
        print('Using saved SVD decomposition')
        U, S, Vh = torch.load(save_path)
        U = U.to(device)
        S = S.to(device)
    else:
        U, S, Vh = torch.svd(x)
        torch.save((U, S, Vh), save_path)
    
    x = U[:,:k] @ torch.diag(S[:k])
    x_s = x[:x_s.shape[0]]
    x_t = x[x_s.shape[0]:]
    x_tv = x_tv[:,:k] # NOT GOOD, but I don't need validation for this experiment
    
    x_s_n = x_s - x_s.mean(0) # centered source
    x_t_n = x_t - x_t.mean(0) # centered target
    x_tv = x_tv - x_t.mean(0)

    x_s_cov = torch.matmul(x_s_n.T, x_s_n) / (x_s_n.shape[0] - 1.)
    x_s_cov = x_s_cov + 0.01 * torch.eye(x_s_cov.shape[0]).to(device)
    x_t_cov = torch.matmul(x_t_n.T, x_t_n) / (x_t_n.shape[0] - 1.)
    x_t_cov = x_t_cov + 0.01 * torch.eye(x_t_cov.shape[0]).to(device)

    x_s_cov_sqrt = torch.tensor(sqrtm(x_s_cov.cpu())).to(device)
    x_s_cov_sqrt_inv = x_s_cov_sqrt.inverse()
    x_s_whitened = torch.matmul(x_s_n, x_s_cov_sqrt_inv.float()) # whiten
    x_t_cov_sqrt = torch.tensor(sqrtm(x_t_cov.cpu())).to(device)
    x_s = torch.matmul(x_s_whitened, x_t_cov_sqrt.float()) # recolor with target variance

    x_tu = x_t_n # centered target
    x_t = x_s[source_features.shape[0]:] # target
    x_s = x_s[:source_features.shape[0]] # source
    
    # target unlabeled, target labeled, source, target validation
    return x_tu, x_t, x_s, x_tv