jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def gen_orthogonal_train_test(data_sampler, n_points, b_size):
    xs = data_sampler.sample_xs(n_points, b_size)
    n_dim = xs.shape[2]
    n_points = min(n_points, n_dim)
    # raise ValueError("number of points should be at most the dimension.")
    xs_train_pre = xs
    xs_test_post = torch.zeros(xs.shape)
    for i in range(n_points):
        xs_test_post_i = xs[:, i : i + 1, :]
        xs_train_pre_i = xs[:, :i, :]
        U, S, Vt = torch.svd(xs_train_pre_i)
        xs_train_pre_i_projection = Vt.transpose(1, 2) @ Vt
        xs_test_post_i_orthogonalized = (
            xs_test_post_i - xs_test_post_i @ xs_train_pre_i_projection
        )
        xs_test_post_i_normalized = (
            xs_test_post_i_orthogonalized
            * xs_test_post_i.norm(dim=2).unsqueeze(2)
            / xs_test_post_i_orthogonalized.norm(dim=2).unsqueeze(2)
        )

        xs_test_post[:, i : i + 1, :] = xs_test_post_i_normalized

    return xs_train_pre, xs_test_post
//  [python]
def gen_orthogonal_train_test(data_sampler, n_points, b_size):
    xs = data_sampler.sample_xs(n_points, b_size)
    n_dim = xs.shape[2]
    n_points = min(n_points, n_dim)
    # raise ValueError("number of points should be at most the dimension.")
    xs_train_pre = xs
    xs_test_post = torch.zeros(xs.shape)
    for i in range(n_points):
        xs_test_post_i = xs[:, i : i + 1, :]
        xs_train_pre_i = xs[:, :i, :]
        _, _, Vt = torch.linalg.svd(xs_train_pre_i, full_matrices=False)
        xs_train_pre_i_projection = Vt.transpose(1, 2) @ Vt
        xs_test_post_i_orthogonalized = (
            xs_test_post_i - xs_test_post_i @ xs_train_pre_i_projection
        )
        xs_test_post_i_normalized = (
            xs_test_post_i_orthogonalized
            * xs_test_post_i.norm(dim=2).unsqueeze(2)
            / xs_test_post_i_orthogonalized.norm(dim=2).unsqueeze(2)
        )

        xs_test_post[:, i : i + 1, :] = xs_test_post_i_normalized

    return xs_train_pre, xs_test_post
//  [python]
def pinverse(inputs: torch.Tensor):
    assert(inputs.ndim >= 2)
    shp = inputs.shape
    U, S, V = torch.svd(inputs.view(-1, shp[-2], shp[-1]))
    S[S < 1e-6] = 0
    S_inv = torch.where(S < 1e-5, torch.zeros_like(S), 1/S)
    pinv = V @ torch.diag_embed(S_inv) @ U.transpose(1,2)
    return pinv.view(shp)
//  [python]
def pinverse(inputs: torch.Tensor):
    assert(inputs.ndim >= 2)
    shp = inputs.shape
    U, S, V = torch.linalg.svd(inputs.view(-1, shp[-2], shp[-1]))
    S[S < 1e-6] = 0
    S_inv = torch.where(S < 1e-5, torch.zeros_like(S), 1/S)
    pinv = V @ torch.diag_embed(S_inv) @ U.transpose(1,2)
    return pinv.view(shp)
//  [python]
def get_svd(_model, toxic_vector, num_mlp_vecs):
    scores = []
    for layer in range(_model.cfg.n_layers):
        mlp_outs = _model.blocks[layer].mlp.W_out
        cos_sims = F.cosine_similarity(
            mlp_outs, toxic_vector.unsqueeze(0), dim=1
        )
        _topk = cos_sims.topk(k=300)
        _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]]
        for x in sorted_scores[:num_mlp_vecs]
    ]
    top_vecs = [x / x.norm() for x in top_vecs]
    _top_vecs = torch.stack(top_vecs)

    svd = torch.svd(_top_vecs.transpose(0, 1))
    return svd, sorted_scores
//  [python]
def get_svd(_model, toxic_vector, num_mlp_vecs):
    scores = []
    for layer in range(_model.cfg.n_layers):
        mlp_outs = _model.blocks[layer].mlp.W_out
        cos_sims = F.cosine_similarity(
            mlp_outs, toxic_vector.unsqueeze(0), dim=1
        )
        _topk = cos_sims.topk(k=300)
        _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]]
        for x in sorted_scores[:num_mlp_vecs]
    ]
    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))
    return svd, sorted_scores
//  [python]
def _svd_weighted_adapter(self, adapters, weights, new_rank, target, target_lora_A, target_lora_B):
        delta_weight = weights[0] * target.get_delta_weight(adapters[0])
        for adapter, weight in zip(adapters[1:], weights[1:]):
            delta_weight += weight * target.get_delta_weight(adapter)
        conv2d = isinstance(target, Conv2d)
        if conv2d:
            conv2d_1x1 = target.weight.size()[2:4] == (1, 1)
            if not conv2d_1x1:
                delta_weight = delta_weight.flatten(start_dim=1)
            else:
                delta_weight = delta_weight.squeeze()
        if target.fan_in_fan_out:
            delta_weight = delta_weight.T

        # based on https://github.com/kohya-ss/sd-scripts/blob/main/networks/svd_merge_lora.py#L114-L131
        U, S, Vh = torch.svd(delta_weight)
        U = U[:, :new_rank]
        S = S[:new_rank]
        U = U @ torch.diag(S)
        Vh = Vh[:new_rank, :]
        dist = torch.cat([U.flatten(), Vh.flatten()])
        hi_val = torch.quantile(dist, CLAMP_QUANTILE)
        low_val = -hi_val
        U = U.clamp(low_val, hi_val)
        Vh = Vh.clamp(low_val, hi_val)
        if conv2d:
            U = U.reshape(target_lora_B.data.shape)
            Vh = Vh.reshape(target_lora_A.data.shape)
        return Vh, U
//  [python]
def _svd_weighted_adapter(self, adapters, weights, new_rank, target, target_lora_A, target_lora_B):
        delta_weight = weights[0] * target.get_delta_weight(adapters[0])
        for adapter, weight in zip(adapters[1:], weights[1:]):
            delta_weight += weight * target.get_delta_weight(adapter)
        conv2d = isinstance(target, Conv2d)
        if conv2d:
            conv2d_1x1 = target.weight.size()[2:4] == (1, 1)
            if not conv2d_1x1:
                delta_weight = delta_weight.flatten(start_dim=1)
            else:
                delta_weight = delta_weight.squeeze()
        if target.fan_in_fan_out:
            delta_weight = delta_weight.T

        # based on https://github.com/kohya-ss/sd-scripts/blob/main/networks/svd_merge_lora.py#L114-L131
        U, S, Vh = torch.linalg.svd(delta_weight)
        U = U[:, :new_rank]
        S = S[:new_rank]
        U = U @ torch.diag(S)
        Vh = Vh[:new_rank, :]
        dist = torch.cat([U.flatten(), Vh.flatten()])
        hi_val = torch.quantile(dist, CLAMP_QUANTILE)
        low_val = -hi_val
        U = U.clamp(low_val, hi_val)
        Vh = Vh.clamp(low_val, hi_val)
        if conv2d:
            U = U.reshape(target_lora_B.data.shape)
            Vh = Vh.reshape(target_lora_A.data.shape)
        return Vh, U
//  [python]
def procrustes_analysis(X0: torch.Tensor,X1: torch.Tensor) -> Dict[str, Any]: # [N,3]
    """Computes 3D similarity between two coordinate systems. 

    Args:
        X0 (torch.Tensor): Center of cameras in coordinate system 0, (B, 3)
        X1 (torch.Tensor_type_): Center of cameras in coordinate system 1, (B, 3)
    """
    # translation
    t0 = X0.mean(dim=0,keepdim=True)
    t1 = X1.mean(dim=0,keepdim=True)
    X0c = X0-t0
    X1c = X1-t1
    # scale
    s0 = (X0c**2).sum(dim=-1).mean().sqrt()
    s1 = (X1c**2).sum(dim=-1).mean().sqrt()
    X0cs = X0c/s0
    X1cs = X1c/s1
    # rotation (use double for SVD, float loses precision)
    U, S, Vh = torch.svd((X0cs.t()@X1cs).cpu().double())
    V = Vh.transpose(-2, -1).conj()
    R = (U@V.t()).float().to(X0.device)
    if R.det()<0: R[2] *= -1
    # align X1 to X0: X1to0 = (X1-t1)/s1@R.t()*s0+t0
    sim3 = edict(t0=t0[0],t1=t1[0],s0=s0,s1=s1,R=R)
    return sim3
//  [python]
def procrustes_analysis(X0: torch.Tensor,X1: torch.Tensor) -> Dict[str, Any]: # [N,3]
    """Computes 3D similarity between two coordinate systems. 

    Args:
        X0 (torch.Tensor): Center of cameras in coordinate system 0, (B, 3)
        X1 (torch.Tensor_type_): Center of cameras in coordinate system 1, (B, 3)
    """
    # translation
    t0 = X0.mean(dim=0,keepdim=True)
    t1 = X1.mean(dim=0,keepdim=True)
    X0c = X0-t0
    X1c = X1-t1
    # scale
    s0 = (X0c**2).sum(dim=-1).mean().sqrt()
    s1 = (X1c**2).sum(dim=-1).mean().sqrt()
    X0cs = X0c/s0
    X1cs = X1c/s1
    # rotation (use double for SVD, float loses precision)
    U, S, Vh = torch.linalg.svd((X0cs.t()@X1cs).cpu().double(), full_matrices=False)
    V = Vh.transpose(-2, -1).conj()
    R = (U@V.t()).float().to(X0.device)
    if R.det()<0: R[2] *= -1
    # align X1 to X0: X1to0 = (X1-t1)/s1@R.t()*s0+t0
    sim3 = edict(t0=t0[0],t1=t1[0],s0=s0,s1=s1,R=R)
    return sim3
//  [python]
def update_model_reduce_layer(self, layer_type, layer_number):
        layer_id = f"{layer_type}_{layer_number}"
        if layer_id in self.modified_layers:
            print(f"Layer {layer_id} has already been modified. Skipping.")
            return False

        for name, module in self.model.named_modules():
            if layer_type in name and str(layer_number) in name:
                print(f"Reconstructing layer: {name}")
                original_dtype = module.weight.dtype
                self.original_weights[name] = module.weight.detach().clone()
                weights = module.weight.double()
                U, S, V = torch.svd(weights)

                # Estimate sigma using the full IQR method
                sigma_estimated_full_iqr = self.estimate_sigma_with_full_iqr(S)

                # Calculate Marchenko-Pastur threshold
                n, m = weights.shape
                mp_threshold_full_iqr = self.marchenko_pastur_threshold(sigma_estimated_full_iqr, n, m)

                # Retain only the singular values above the MP threshold
                S_reduced = torch.zeros_like(S)
                k = (S > mp_threshold_full_iqr).sum().item()
                S_reduced[:k] = S[:k]
                print(f"Reduced from {S.shape} to {k}")

                # Reconstruct the matrix using the thresholded singular values
                reconstructed_weights = U @ torch.diag(S_reduced) @ V.t()
                reconstructed_weights = reconstructed_weights.to(original_dtype)
                module.weight = torch.nn.Parameter(reconstructed_weights)
                self.modified_layers.add(layer_id)
                return True
//  [python]
def update_model_reduce_layer(self, layer_type, layer_number):
        layer_id = f"{layer_type}_{layer_number}"
        if layer_id in self.modified_layers:
            print(f"Layer {layer_id} has already been modified. Skipping.")
            return False

        for name, module in self.model.named_modules():
            if layer_type in name and str(layer_number) in name:
                print(f"Reconstructing layer: {name}")
                original_dtype = module.weight.dtype
                self.original_weights[name] = module.weight.detach().clone()
                weights = module.weight.double()
                U, S, V = torch.linalg.svd(weights, full_matrices=False)

                # Estimate sigma using the full IQR method
                sigma_estimated_full_iqr = self.estimate_sigma_with_full_iqr(S)

                # Calculate Marchenko-Pastur threshold
                n, m = weights.shape
                mp_threshold_full_iqr = self.marchenko_pastur_threshold(sigma_estimated_full_iqr, n, m)

                # Retain only the singular values above the MP threshold
                S_reduced = torch.zeros_like(S)
                k = (S > mp_threshold_full_iqr).sum().item()
                S_reduced[:k] = S[:k]
                print(f"Reduced from {S.shape} to {k}")

                # Reconstruct the matrix using the thresholded singular values
                reconstructed_weights = U @ torch.diag(S_reduced) @ V
                reconstructed_weights = reconstructed_weights.to(original_dtype)
                module.weight = torch.nn.Parameter(reconstructed_weights)
                self.modified_layers.add(layer_id)
                return True
//  [python]
def rrr(X, Y, rank=8):
    B_ols = ols(X, Y) # if infeasible use GD (numpy CG)
    U, S, V = t.svd(Y.T @ X @ B_ols)
    W = V[0:rank].T

    B_rr = B_ols @ W @ W.T
    L = B_rr#.T
    return L
//  [python]
def rrr(X, Y, rank=8):
    B_ols = ols(X, Y) # if infeasible use GD (numpy CG)
    U, S, V = torch.linalg.svd(Y.T @ X @ B_ols)
    W = V[0:rank].T

    B_rr = B_ols @ W @ W.T
    L = B_rr#.T
    return L
//  [python]
def test_invariance_error_spectral_decompositions(self, device, dtype):
        make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=True)
        A = make_arg((3, 3))
        with self.assertRaisesRegex(RuntimeError, "ill-defined"):
            U, _, Vh = torch.svd(A, some=False)
            (U + Vh).sum().backward()

        A = make_arg((3, 3))
        with self.assertRaisesRegex(RuntimeError, "ill-defined"):
            V = torch.eig(A).eigenvectors
            V.sum().backward()

        A = make_arg((3, 3))
        A = A + A.mH
        with self.assertRaisesRegex(RuntimeError, "ill-defined"):
            Q = torch.eigh(A).eigenvectors
            Q.sum().backward()
//  [python]
def test_invariance_error_spectral_decompositions(self, device, dtype):
        make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=True)
        A = make_arg((3, 3))
        with self.assertRaisesRegex(RuntimeError, "ill-defined"):
            U, _, Vh = torch.linalg.svd(A, full_matrices=False)
            (U + Vh).sum().backward()

        A = make_arg((3, 3))
        with self.assertRaisesRegex(RuntimeError, "ill-defined"):
            V = torch.linalg.eig(A).eigenvectors
            V.sum().backward()

        A = make_arg((3, 3))
        A = A + A.mH
        with self.assertRaisesRegex(RuntimeError, "ill-defined"):
            Q = torch.linalg.eigh(A).eigenvectors
            Q.sum().backward()
//  [python]
def test_svd_memory_allocation(self, device, dtype):
        # test for https://github.com/pytorch/pytorch/issues/61949
        # the problem was that tensors of incorrect size were allocated and then narrowed
        m = 3
        n = 2**20
        a = make_tensor((m, n), dtype=dtype, device=device)
        # the following should run without errors
        S = torch.linalg.svdvals(a)
        result = torch.svd(a, some=False)
        self.assertEqual(result.S, S)
//  [python]
def test_svd_memory_allocation(self, device, dtype):
        # test for https://github.com/pytorch/pytorch/issues/61949
        # the problem was that tensors of incorrect size were allocated and then narrowed
        m = 3
        n = 2**20
        a = make_tensor((m, n), dtype=dtype, device=device)
        # the following should run without errors
        S = torch.linalg.svdvals(a)
        result = torch.linalg.svd(a, full_matrices=False)
        self.assertEqual(result.S, S)
//  [python]
def truncated_SVD(self, traj, k=None, full_matrices=False):
        r"""Truncated Singular Value Decomposition

        Args:
            traj (torch.Tensor): The trajectory to be decomposed
            k (int): The number of singular values and vectors to be computed
            full_matrices (bool): Whether to compute full-sized matrices

        Returns:
            U_trunc (torch.Tensor): The truncated left singular vectors
            S_trunc (torch.Tensor): The truncated singular values
            Vt_trunc (torch.Tensor): The truncated right singular vectors
        """

        assert traj.size(2) == self.dim  # NTC
        k = self.k if k is None else k

        # Singular Value Decomposition
        M = traj.reshape(-1, traj.size(1) * self.dim).T
        U, S, Vt = torch.svd(M, some=full_matrices)

        # Truncated SVD
        U_trunc, S_trunc, Vt_trunc = U[:, :k], S[:k], Vt[:k, :]
        return U_trunc, S_trunc, Vt_trunc.T
//  [python]
def truncated_SVD(self, traj, k=None, full_matrices=False):
        r"""Truncated Singular Value Decomposition

        Args:
            traj (torch.Tensor): The trajectory to be decomposed
            k (int): The number of singular values and vectors to be computed
            full_matrices (bool): Whether to compute full-sized matrices

        Returns:
            U_trunc (torch.Tensor): The truncated left singular vectors
            S_trunc (torch.Tensor): The truncated singular values
            Vt_trunc (torch.Tensor): The truncated right singular vectors
        """

        assert traj.size(2) == self.dim  # NTC
        k = self.k if k is None else k

        # Singular Value Decomposition
        M = traj.reshape(-1, traj.size(1) * self.dim).T
        U, S, Vt = torch.linalg.svd(M, full_matrices=full_matrices)

        # Truncated SVD
        U_trunc, S_trunc, Vt_trunc = U[:, :k], S[:k], Vt[:k, :]
        return U_trunc, S_trunc, Vt_trunc.T
//  [python]
def test_svd(self, device, dtype):
        # tests linalg.svd, svd, linalg.svdvals
        make_arg = partial(make_tensor, dtype=dtype, device=device)

        backends = ["default"]

        if torch.device(device).type == 'cuda':
            if torch.cuda.has_magma:
                backends.append("magma")
            if has_cusolver():
                backends.append("cusolver")

        ns = (5, 2, 0)
        batches = ((), (0,), (1,), (2,), (2, 1), (0, 2))

        for backend in backends:
            torch.backends.cuda.preferred_linalg_library(backend)

            for batch, m, n in product(batches, ns, ns):
                shape = batch + (m, n)
                k = min(m, n)
                A = make_arg(shape)
                U, S, Vh = torch.svd(A, full_matrices=False)
                self.assertEqual((U @ S.to(A.dtype).diag_embed()) @ Vh, A)

                U_f, S_f, Vh_f = torch.svd(A, full_matrices=True)
                self.assertEqual(S_f, S)
                self.assertEqual((U_f[..., :k] @ S_f.to(A.dtype).diag_embed()) @ Vh_f[..., :k, :], A)

                S_s = torch.svd(A, compute_uv=False).S
                self.assertEqual(S_s, S)

                U, S, V = torch.svd(A, some=True)
                self.assertEqual((U @ S.to(A.dtype).diag_embed()) @ V.mH, A)

                U_f, S_f, V_f = torch.svd(A, some=False)
                self.assertEqual(S_f, S)
                self.assertEqual((U_f[..., :k] @ S_f.to(A.dtype).diag_embed()) @ V_f[..., :k].mH, A)

                S_s = torch.svd(A, compute_uv=False).S
                self.assertEqual(S_s, S)
//  [python]
def test_svd(self, device, dtype):
        # tests linalg.svd, svd, linalg.svdvals
        make_arg = partial(make_tensor, dtype=dtype, device=device)

        backends = ["default"]

        if torch.device(device).type == 'cuda':
            if torch.cuda.has_magma:
                backends.append("magma")
            if has_cusolver():
                backends.append("cusolver")

        ns = (5, 2, 0)
        batches = ((), (0,), (1,), (2,), (2, 1), (0, 2))

        for backend in backends:
            torch.backends.cuda.preferred_linalg_library(backend)

            for batch, m, n in product(batches, ns, ns):
                shape = batch + (m, n)
                k = min(m, n)
                A = make_arg(shape)
                U, S, Vh = torch.linalg.svd(A, full_matrices=False)
                self.assertEqual((U @ S.to(A.dtype).diag_embed()) @ Vh, A)

                U_f, S_f, Vh_f = torch.linalg.svd(A, full_matrices=True)
                self.assertEqual(S_f, S)
                self.assertEqual((U_f[..., :k] @ S_f.to(A.dtype).diag_embed()) @ Vh_f[..., :k, :], A)

                S_s = torch.linalg.svdvals(A)
                self.assertEqual(S_s, S)

                U, S, V = torch.svd(A, some=True)
                self.assertEqual((U @ S.to(A.dtype).diag_embed()) @ V.mH, A)

                U_f, S_f, V_f = torch.svd(A, some=False)
                self.assertEqual(S_f, S)
                self.assertEqual((U_f[..., :k] @ S_f.to(A.dtype).diag_embed()) @ V_f[..., :k].mH, A)

                S_s = torch.svd(A, compute_uv=False).S
                self.assertEqual(S_s, S)
//  [python]
def batch_compute_similarity_transform_torch(S1, S2):
    '''
    Computes a similarity transform (sR, t) that takes
    a set of 3D points S1 (3 x N) closest to a set of 3D points S2,
    where R is an 3x3 rotation matrix, t 3x1 translation, s scale.
    i.e. solves the orthogonal Procrutes problem.
    '''
    transposed = False
    if S1.shape[1] != 3 and S1.shape[1] != 2:
        S1 = S1.permute(0,2,1)
        S2 = S2.permute(0,2,1)
        transposed = True
    assert(S2.shape[1] == S1.shape[1])

    # 1. Remove mean.
    mu1 = S1.mean(axis=-1, keepdims=True)
    mu2 = S2.mean(axis=-1, keepdims=True)
    X1 = S1 - mu1
    X2 = S2 - mu2
    # 2. Compute variance of X1 used for scale.
    var1 = torch.sum(X1**2, dim=1).sum(dim=1)
    # 3. The outer product of X1 and X2.
    K = X1.bmm(X2.permute(0,2,1))
    # 4. Solution that Maximizes trace(R'K) is R=U*V', where U, V are
    # singular vectors of K.
    U, s, Vh = torch.svd(K)
    V = Vh.transpose(-2,-1)
    # Construct Z that fixes the orientation of R to get det(R)=1.
    Z = torch.eye(U.shape[1], device=S1.device).unsqueeze(0)
    Z = Z.repeat(U.shape[0],1,1)
    Z[:,-1, -1] *= torch.sign(torch.det(U.bmm(V.permute(0,2,1))))
    # Construct R.
    R = V.bmm(Z.bmm(U.permute(0,2,1)))
    # 5. Recover scale.
    scale = torch.cat([torch.trace(x).unsqueeze(0) for x in R.bmm(K)]) / var1
    # 6. Recover translation.
    t = mu2 - (scale.unsqueeze(-1).unsqueeze(-1) * (R.bmm(mu1)))
    # 7. Error:
    S1_hat = scale.unsqueeze(-1).unsqueeze(-1) * R.bmm(S1) + t
    if transposed:
        S1_hat = S1_hat.permute(0,2,1)

    return S1_hat, (scale, R, t)
//  [python]
def batch_compute_similarity_transform_torch(S1, S2):
    '''
    Computes a similarity transform (sR, t) that takes
    a set of 3D points S1 (3 x N) closest to a set of 3D points S2,
    where R is an 3x3 rotation matrix, t 3x1 translation, s scale.
    i.e. solves the orthogonal Procrutes problem.
    '''
    transposed = False
    if S1.shape[1] != 3 and S1.shape[1] != 2:
        S1 = S1.permute(0,2,1)
        S2 = S2.permute(0,2,1)
        transposed = True
    assert(S2.shape[1] == S1.shape[1])

    # 1. Remove mean.
    mu1 = S1.mean(axis=-1, keepdims=True)
    mu2 = S2.mean(axis=-1, keepdims=True)
    X1 = S1 - mu1
    X2 = S2 - mu2
    # 2. Compute variance of X1 used for scale.
    var1 = torch.sum(X1**2, dim=1).sum(dim=1)
    # 3. The outer product of X1 and X2.
    K = X1.bmm(X2.permute(0,2,1))
    # 4. Solution that Maximizes trace(R'K) is R=U*V', where U, V are
    # singular vectors of K.
    U, s, Vh = torch.linalg.svd(K)
    V = Vh.transpose(-2,-1)
    # Construct Z that fixes the orientation of R to get det(R)=1.
    Z = torch.eye(U.shape[1], device=S1.device).unsqueeze(0)
    Z = Z.repeat(U.shape[0],1,1)
    Z[:,-1, -1] *= torch.sign(torch.det(U.bmm(V.permute(0,2,1))))
    # Construct R.
    R = V.bmm(Z.bmm(U.permute(0,2,1)))
    # 5. Recover scale.
    scale = torch.cat([torch.trace(x).unsqueeze(0) for x in R.bmm(K)]) / var1
    # 6. Recover translation.
    t = mu2 - (scale.unsqueeze(-1).unsqueeze(-1) * (R.bmm(mu1)))
    # 7. Error:
    S1_hat = scale.unsqueeze(-1).unsqueeze(-1) * R.bmm(S1) + t
    if transposed:
        S1_hat = S1_hat.permute(0,2,1)

    return S1_hat, (scale, R, t)
//  [python]
def fit(self, x):
        n = x.shape[0]
        x = x-torch.mean(x, dim=0)
        covariance_matrix = 1 / n * (x.T@x)
        '''svd'''
        u, s, vT = torch.svd(covariance_matrix)
        s_argsort = torch.argsort(-s)
        fit_x = u[:, [*s_argsort[0:self.principal_components]]]
        '''evd'''
        # eigen_values, eigen_vector = torch.linalg.eig(covariance_matrix)
        # eigen_values_argsort = torch.argsort(-eigen_values.real)
        # fit_x = eigen_vector[:, [*eigen_values_argsort[0:self.principal_components]]].real
        return fit_x
//  [python]
def fit(self, x):
        n = x.shape[0]
        x = x-torch.mean(x, dim=0)
        covariance_matrix = 1 / n * (x.T@x)
        '''svd'''
        u, s, vT = torch.linalg.svd(covariance_matrix)
        s_argsort = torch.argsort(-s)
        fit_x = u[:, [*s_argsort[0:self.principal_components]]]
        '''evd'''
        # eigen_values, eigen_vector = torch.linalg.eig(covariance_matrix)
        # eigen_values_argsort = torch.argsort(-eigen_values.real)
        # fit_x = eigen_vector[:, [*eigen_values_argsort[0:self.principal_components]]].real
        return fit_x
//  [python]
def _svd_aten(
    A: TensorLikeType, *, full_matrices: bool
) -> Tuple[Tensor, Tensor, Tensor]:
    return torch.svd(A, full_matrices=full_matrices)
//  [python]
def _svd_aten(
    A: TensorLikeType, *, full_matrices: bool
) -> Tuple[Tensor, Tensor, Tensor]:
    return torch.linalg.svd(A, full_matrices=full_matrices)
//  [python]
def estimate_non_minimal_model(self, pts, T1, T2t, weights=None):  # x1 y1 x2 y2
        """
        Using 8 points and singularity constraint to estimate Fundamental matrix.
        """
        dev = pts.device
        # get the points
        B, N, _ = pts.shape
        #matches, transform1, transform2 = self.normalize(pts)
        pts1 = pts[:, :, 0:2]
        pts2 = pts[:, :, 2:4]
        x1, y1 = pts1[:, :, 0], pts1[:, :, 1]
        x2, y2 = pts2[:, :, 0], pts2[:, :, 1]

        a_89 = torch.ones(x1.shape, device=dev)

        # construct the A matrix, A F = 0. 8 equations for 9 variables, solution is linear subspace o dimensionality of 2.
        #print(weights)
        if weights is not None:
            A = weights.unsqueeze(-1) * torch.stack((x1 * x2, x2 * y1, x2, y2 * x1, y2 * y1, y2, x1, y1, a_89), dim=-1)#weights.unsqueeze(-1) *
        else:
            A = torch.stack((x1 * x2, x2 * y1, x2, y2 * x1, y2 * y1, y2, x1, y1, a_89), dim=-1)

        # solve null space of A to get F
        try:
            U, S, V = torch.svd(A.transpose(-1, -2)@A)#, full_matrices=False)  # eigenvalues in increasing order
        except:
            print()
        null_space = V[:, -1:, :].transpose(-1, -2).float().clone()  # the last four rows

        # with the singularity constraint, use the last two singular vectors as a basis of the space
        F = null_space[:, :, 0].view(-1, 3, 3)

        if T1 is not None:
            for i in range(F.shape[0]):
                F[i, :, :] = torch.mm(T2t[i, :, :], torch.mm(F[i, :, :].clone(), T1[i, :, :]))

        return F, None
//  [python]
def estimate_non_minimal_model(self, pts, T1, T2t, weights=None):  # x1 y1 x2 y2
        """
        Using 8 points and singularity constraint to estimate Fundamental matrix.
        """
        dev = pts.device
        # get the points
        B, N, _ = pts.shape
        #matches, transform1, transform2 = self.normalize(pts)
        pts1 = pts[:, :, 0:2]
        pts2 = pts[:, :, 2:4]
        x1, y1 = pts1[:, :, 0], pts1[:, :, 1]
        x2, y2 = pts2[:, :, 0], pts2[:, :, 1]

        a_89 = torch.ones(x1.shape, device=dev)

        # construct the A matrix, A F = 0. 8 equations for 9 variables, solution is linear subspace o dimensionality of 2.
        #print(weights)
        if weights is not None:
            A = weights.unsqueeze(-1) * torch.stack((x1 * x2, x2 * y1, x2, y2 * x1, y2 * y1, y2, x1, y1, a_89), dim=-1)#weights.unsqueeze(-1) *
        else:
            A = torch.stack((x1 * x2, x2 * y1, x2, y2 * x1, y2 * y1, y2, x1, y1, a_89), dim=-1)

        # solve null space of A to get F
        try:
            _, _, v = torch.linalg.svd(A.transpose(-1, -2)@A)#, full_matrices=False)  # eigenvalues in increasing order
        except:
            print()
        null_space = v[:, -1:, :].transpose(-1, -2).float().clone()  # the last four rows

        # with the singularity constraint, use the last two singular vectors as a basis of the space
        F = null_space[:, :, 0].view(-1, 3, 3)

        if T1 is not None:
            for i in range(F.shape[0]):
                F[i, :, :] = torch.mm(T2t[i, :, :], torch.mm(F[i, :, :].clone(), T1[i, :, :]))

        return F, None
//  [python]
def extract_lora(diff, rank):
    conv2d = (len(diff.shape) == 4)
    kernel_size = None if not conv2d else diff.size()[2:4]
    conv2d_3x3 = conv2d and kernel_size != (1, 1)
    out_dim, in_dim = diff.size()[0:2]
    rank = min(rank, in_dim, out_dim)

    if conv2d:
        if conv2d_3x3:
            diff = diff.flatten(start_dim=1)
        else:
            diff = diff.squeeze()


    U, S, Vh = torch.svd(diff.float())
    U = U[:, :rank]
    S = S[:rank]
    U = U @ torch.diag(S)
    Vh = Vh[:rank, :]

    dist = torch.cat([U.flatten(), Vh.flatten()])
    hi_val = torch.quantile(dist, CLAMP_QUANTILE)
    low_val = -hi_val

    U = U.clamp(low_val, hi_val)
    Vh = Vh.clamp(low_val, hi_val)
    if conv2d:
        U = U.reshape(out_dim, rank, 1, 1)
        Vh = Vh.reshape(rank, in_dim, kernel_size[0], kernel_size[1])
    return (U, Vh)
//  [python]
def extract_lora(diff, rank):
    conv2d = (len(diff.shape) == 4)
    kernel_size = None if not conv2d else diff.size()[2:4]
    conv2d_3x3 = conv2d and kernel_size != (1, 1)
    out_dim, in_dim = diff.size()[0:2]
    rank = min(rank, in_dim, out_dim)

    if conv2d:
        if conv2d_3x3:
            diff = diff.flatten(start_dim=1)
        else:
            diff = diff.squeeze()


    U, S, Vh = torch.linalg.svd(diff.float())
    U = U[:, :rank]
    S = S[:rank]
    U = U @ torch.diag(S)
    Vh = Vh[:rank, :]

    dist = torch.cat([U.flatten(), Vh.flatten()])
    hi_val = torch.quantile(dist, CLAMP_QUANTILE)
    low_val = -hi_val

    U = U.clamp(low_val, hi_val)
    Vh = Vh.clamp(low_val, hi_val)
    if conv2d:
        U = U.reshape(out_dim, rank, 1, 1)
        Vh = Vh.reshape(rank, in_dim, kernel_size[0], kernel_size[1])
    return (U, Vh)
//  [python]
def _low_rank_decomposition(weight, reduced_rank=32):
    """
    :param weight: The matrix to decompose, of shape (H, W) :param reduced_rank: the final rank :return:
    """
    matrix_dimension = len(weight.size())
    if matrix_dimension != 2:
        raise ValueError(f"Only support 2D matrix, but your input has {matrix_dimension} dimensions.")

    # Use SVD to decompose a matrix, default full_matrices is False to save parameters
    U, S, Vh = torch.svd(weight)

    L = U @ (torch.sqrt(torch.diag(S)[:, 0:reduced_rank]))
    R = torch.sqrt(torch.diag(S)[0:reduced_rank, :]) @ Vh

    return {"L": L, "R": R, "U": U, "S": S, "Vh": Vh, "reduced_rank": reduced_rank}
//  [python]
def _low_rank_decomposition(weight, reduced_rank=32):
    """
    :param weight: The matrix to decompose, of shape (H, W) :param reduced_rank: the final rank :return:
    """
    matrix_dimension = len(weight.size())
    if matrix_dimension != 2:
        raise ValueError(f"Only support 2D matrix, but your input has {matrix_dimension} dimensions.")

    # Use SVD to decompose a matrix, default full_matrices is False to save parameters
    U, S, Vh = torch.linalg.svd(weight, full_matrices=False)

    L = U @ (torch.sqrt(torch.diag(S)[:, 0:reduced_rank]))
    R = torch.sqrt(torch.diag(S)[0:reduced_rank, :]) @ Vh

    return {"L": L, "R": R, "U": U, "S": S, "Vh": Vh, "reduced_rank": reduced_rank}
//  [python]
def __init__(self, num_nodes, max_num_edges, dim, static_feat, device='cuda', alpha1=0.1, alpha2=2.0,
                 self_loops=True,num_eofs=100):
        super().__init__()
        if static_feat is None:
            raise ValueError("Please give static node features (e.g. part of the timeseries)")
        self.num_nodes = num_nodes

        A, Lh, E = torch.svd(torch.from_numpy(static_feat).T)
        A = A[:, :len(Lh)] #trim the time field (in case time > space)
        A = A[:, :len(Lh)] #trim the time field (in case time > space)
        PCs = A*Lh
        select_neofs=22
        self.select_neofs=num_eofs

        # normalize time series and scale in singular values to retain variance
        self.eofs = E[:self.select_neofs,:]
        n=torch.var(self.eofs,axis=1)
        self.eofs=self.eofs/torch.unsqueeze(n,dim=1)

        self.PCS=PCs[:,:self.select_neofs]
        self.PCS = self.PCS*torch.unsqueeze(n,dim=1).T*Lh[:self.select_neofs]
        self.eig_vals=Lh

        A, Lh, E = torch.svd(torch.from_numpy(static_feat).T)
        A = A[:, :len(Lh)] #trim the time field (in case time > space)
        A = A[:, :len(Lh)] #trim the time field (in case time > space)

        self.var_explained=(Lh*Lh)/(torch.sum(Lh*Lh))
        tau0_data = (self.PCS[:,:self.select_neofs].T)

        tlag = 1
        self.tlag = tlag
        x0=tau0_data[:,0:-tlag]
        x1=tau0_data[:,tlag:]

        #covariance structrure
        self.static_feat = x0 if isinstance(x0, torch.Tensor) else torch.from_numpy(x0)
        self.static_feat = self.static_feat.to(torch.float64).to(device)

        #lagged feature
        self.static_feat1 = x1 if isinstance(x1, torch.Tensor) else torch.from_numpy(x1)
        self.static_feat1 = self.static_feat1.to(torch.float64).to(device)

        self.device = device
        self.dim = dim
        self.alpha1 = alpha1
        self.alpha2 = alpha2
        self.num_edges = max_num_edges
        self.self_loops = self_loops
        self.diag = torch.eye(self.num_nodes).bool().to(device)
//  [python]
def __init__(self, num_nodes, max_num_edges, dim, static_feat, device='cuda', alpha1=0.1, alpha2=2.0,
                 self_loops=True,num_eofs=100):
        super().__init__()
        if static_feat is None:
            raise ValueError("Please give static node features (e.g. part of the timeseries)")
        self.num_nodes = num_nodes

        A, Lh, E = torch.linalg.svd(torch.from_numpy(static_feat).T)
        A = A[:, :len(Lh)] #trim the time field (in case time > space)
        A = A[:, :len(Lh)] #trim the time field (in case time > space)
        PCs = A*Lh
        select_neofs=22
        self.select_neofs=num_eofs

        # normalize time series and scale in singular values to retain variance
        self.eofs = E[:self.select_neofs,:]
        n=torch.var(self.eofs,axis=1)
        self.eofs=self.eofs/torch.unsqueeze(n,dim=1)

        self.PCS=PCs[:,:self.select_neofs]
        self.PCS = self.PCS*torch.unsqueeze(n,dim=1).T*Lh[:self.select_neofs]
        self.eig_vals=Lh

        A, Lh, E = torch.linalg.svd(torch.from_numpy(static_feat).T)
        A = A[:, :len(Lh)] #trim the time field (in case time > space)
        A = A[:, :len(Lh)] #trim the time field (in case time > space)

        self.var_explained=(Lh*Lh)/(torch.sum(Lh*Lh))
        tau0_data = (self.PCS[:,:self.select_neofs].T)

        tlag = 1
        self.tlag = tlag
        x0=tau0_data[:,0:-tlag]
        x1=tau0_data[:,tlag:]

        #covariance structrure
        self.static_feat = x0 if isinstance(x0, torch.Tensor) else torch.from_numpy(x0)
        self.static_feat = self.static_feat.to(torch.float64).to(device)

        #lagged feature
        self.static_feat1 = x1 if isinstance(x1, torch.Tensor) else torch.from_numpy(x1)
        self.static_feat1 = self.static_feat1.to(torch.float64).to(device)

        self.device = device
        self.dim = dim
        self.alpha1 = alpha1
        self.alpha2 = alpha2
        self.num_edges = max_num_edges
        self.self_loops = self_loops
        self.diag = torch.eye(self.num_nodes).bool().to(device)
//  [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.float().to(device)
    out_ch, in_ch, kernel_size, _ = weight.shape
    
    U, S, Vh = torch.linalg.svd(weight.reshape(out_ch, -1))
    
    if 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.cpu(), 'full'
    
    U = U[:, :lora_rank]
    S = S[:lora_rank]
    U = U @ torch.diag(S)
    Vh = Vh[:lora_rank, :]
    
    diff = (weight - (U @ Vh).reshape(out_ch, in_ch, kernel_size, kernel_size)).detach().to(dtype=weight.dtype).cpu()
    extract_weight_A = Vh.reshape(lora_rank, in_ch, kernel_size, kernel_size).detach().to(dtype=weight.dtype).cpu()
    extract_weight_B = U.reshape(out_ch, lora_rank, 1, 1).detach().to(dtype=weight.dtype).cpu()
    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.float().to(device)
    out_ch, in_ch, kernel_size, _ = weight.shape
    
    U, S, Vh = linalg.svd(weight.reshape(out_ch, -1))
    
    if 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.cpu(), 'full'
    
    U = U[:, :lora_rank]
    S = S[:lora_rank]
    U = U @ torch.diag(S)
    Vh = Vh[:lora_rank, :]
    
    diff = (weight - (U @ Vh).reshape(out_ch, in_ch, kernel_size, kernel_size)).detach().to(dtype=weight.dtype).cpu()
    extract_weight_A = Vh.reshape(lora_rank, in_ch, kernel_size, kernel_size).detach().to(dtype=weight.dtype).cpu()
    extract_weight_B = U.reshape(out_ch, lora_rank, 1, 1).detach().to(dtype=weight.dtype).cpu()
    del U, S, Vh, weight
    return (extract_weight_A, extract_weight_B, diff), 'low rank'
//  [python]
def estimate_minimal_model(self, pts, weights=None):  # x1 y1 x2 y2
        """
        using 7 points and singularity constraint to estimate Fundamental matrix.
        """

        dev = pts.device
        # get the points
        B, N, _ = pts.shape
        pts1 = pts[:, :, 0:2]
        pts2 = pts[:, :, 2:4]
        x1, y1 = pts1[:, :, 0], pts1[:, :, 1]
        x2, y2 = pts2[:, :, 0], pts2[:, :, 1]

        a_79 = torch.ones(x1.shape, device=dev)

        # construct the A matrix, A F = 0. 7 equations for 9 variables,
        # solution is linear subspace o dimensionality of 2.
        A = torch.stack((x1 * x2, x2 * y1, x2, y2 * x1, y2 * y1, y2, x1, y1, a_79), dim=-1)

        # solve null space of A to get F
        U, S, V = torch.svd(A.transpose(-1, -2)@A)  # eigenvalues in increasing order
        null_space = V[:, -2:, :].transpose(-1, -2).float()  # the last two rows

        # with the singularity constraint, use the last two singular vectors as a basis of the space
        F1 = null_space[:, :, 0].view(-1, 1, 3, 3)
        F2 = null_space[:, :, 1].view(-1, 1, 3, 3)

        # use the two bases, we can have an arbitrary F mat
        # lambda, 1-lambda, det(F) = det(lambda*F1, (1-lambda)*F2) = lambda(F1-F2)+F2 to find lambda
        # c-polynomial coefficients. det(F) = c[0]*lambda^3 + c[1]*lambda^2  + c[2]*lambda + c[3]= 0
        c = self.coeff(F1, F2).squeeze()

        # Find the roots of the polynomial c by the companion matrix method
        compmat = torch.zeros((c.shape[0], 4, 4), dtype=c.dtype, device=dev)
        compmat[:, 1, 0] = 1.
        compmat[:, 2, 1] = 1.
        compmat[:, 3, 2] = 1.
        compmat[..., 2] = -c
        vv = torch.linalg.eigvals(compmat)

        roots = vv.real
        s = F1[..., 2, 2] * roots + F2[..., 2, 2]
        valid_mask = (s > 1e-10) & (vv.imag.abs() < 1e-9)
        mu = 1.0 / s
        lambda_ = roots * mu
        Fs = F1 * lambda_.view(-1, 4, 1, 1) + F2 * mu.view(-1, 4, 1, 1)
        Fs = Fs[valid_mask].view(-1, 3, 3)
        pl_models = []
        for p, F in zip(pts, Fs):
            ransac_opt = {
                'max_iterations': 1,
                'min_iterations': 1,
                'max_epipolar_error': 1000,
                'progressive_sampling': False
            }
            p1, p2 = p[:, :2].cpu().numpy(), p[:, 2:].cpu().numpy()
            model, stats = poselib.estimate_fundamental(p1, p2, ransac_opt)
            pl_models.append(model)
            inl = stats["inliers"]
        Fs = torch.tensor(np.stack(pl_models, axis=0), dtype=pts.dtype, device=pts.device)
        return Fs, None
//  [python]
def estimate_minimal_model(self, pts, weights=None):  # x1 y1 x2 y2
        """
        using 7 points and singularity constraint to estimate Fundamental matrix.
        """

        dev = pts.device
        # get the points
        B, N, _ = pts.shape
        pts1 = pts[:, :, 0:2]
        pts2 = pts[:, :, 2:4]
        x1, y1 = pts1[:, :, 0], pts1[:, :, 1]
        x2, y2 = pts2[:, :, 0], pts2[:, :, 1]

        a_79 = torch.ones(x1.shape, device=dev)

        # construct the A matrix, A F = 0. 7 equations for 9 variables,
        # solution is linear subspace o dimensionality of 2.
        A = torch.stack((x1 * x2, x2 * y1, x2, y2 * x1, y2 * y1, y2, x1, y1, a_79), dim=-1)

        # solve null space of A to get F
        _, _, v = torch.linalg.svd(A.transpose(-1, -2)@A)  # eigenvalues in increasing order
        null_space = v[:, -2:, :].transpose(-1, -2).float()  # the last two rows

        # with the singularity constraint, use the last two singular vectors as a basis of the space
        F1 = null_space[:, :, 0].view(-1, 1, 3, 3)
        F2 = null_space[:, :, 1].view(-1, 1, 3, 3)

        # use the two bases, we can have an arbitrary F mat
        # lambda, 1-lambda, det(F) = det(lambda*F1, (1-lambda)*F2) = lambda(F1-F2)+F2 to find lambda
        # c-polynomial coefficients. det(F) = c[0]*lambda^3 + c[1]*lambda^2  + c[2]*lambda + c[3]= 0
        c = self.coeff(F1, F2).squeeze()

        # Find the roots of the polynomial c by the companion matrix method
        compmat = torch.zeros((c.shape[0], 4, 4), dtype=c.dtype, device=dev)
        compmat[:, 1, 0] = 1.
        compmat[:, 2, 1] = 1.
        compmat[:, 3, 2] = 1.
        compmat[..., 2] = -c
        vv = torch.linalg.eigvals(compmat)

        roots = vv.real
        s = F1[..., 2, 2] * roots + F2[..., 2, 2]
        valid_mask = (s > 1e-10) & (vv.imag.abs() < 1e-9)
        mu = 1.0 / s
        lambda_ = roots * mu
        Fs = F1 * lambda_.view(-1, 4, 1, 1) + F2 * mu.view(-1, 4, 1, 1)
        Fs = Fs[valid_mask].view(-1, 3, 3)
        pl_models = []
        for p, F in zip(pts, Fs):
            ransac_opt = {
                'max_iterations': 1,
                'min_iterations': 1,
                'max_epipolar_error': 1000,
                'progressive_sampling': False
            }
            p1, p2 = p[:, :2].cpu().numpy(), p[:, 2:].cpu().numpy()
            model, stats = poselib.estimate_fundamental(p1, p2, ransac_opt)
            pl_models.append(model)
            inl = stats["inliers"]
        Fs = torch.tensor(np.stack(pl_models, axis=0), dtype=pts.dtype, device=pts.device)
        return Fs, None
//  [python]
def extract_linear(
    weight: Union[torch.Tensor, nn.Parameter],
    mode = 'fixed',
    mode_param = 0,
    device = 'cpu',
) -> Tuple[nn.Parameter, nn.Parameter]:
    weight = weight.float().to(device)
    out_ch, in_ch = weight.shape
    
    U, S, Vh = torch.linalg.svd(weight)
    
    if 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.cpu(), 'full'
    
    U = U[:, :lora_rank]
    S = S[:lora_rank]
    U = U @ torch.diag(S)
    Vh = Vh[:lora_rank, :]
    
    diff = (weight - U @ Vh).detach().to(dtype=weight.dtype).cpu()
    extract_weight_A = Vh.reshape(lora_rank, in_ch).detach().to(dtype=weight.dtype).cpu()
    extract_weight_B = U.reshape(out_ch, lora_rank).detach().to(dtype=weight.dtype).cpu()
    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.float().to(device)
    out_ch, in_ch = weight.shape
    
    U, S, Vh = linalg.svd(weight)
    
    if 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.cpu(), 'full'
    
    U = U[:, :lora_rank]
    S = S[:lora_rank]
    U = U @ torch.diag(S)
    Vh = Vh[:lora_rank, :]
    
    diff = (weight - U @ Vh).detach().to(dtype=weight.dtype).cpu()
    extract_weight_A = Vh.reshape(lora_rank, in_ch).detach().to(dtype=weight.dtype).cpu()
    extract_weight_B = U.reshape(out_ch, lora_rank).detach().to(dtype=weight.dtype).cpu()
    del U, S, Vh, weight
    return (extract_weight_A, extract_weight_B, diff), 'low rank'
//  [python]
def convert_pca(model, dim, device, embed_dict):
    """
    Perform SVD to reduce the dimensionality of embeddings through the final weight matrix in 
    a resnet model
    Inputs:
    - model: resnet model
    - dim: dimension to reduce to
    - embed_dict: {'embed': array-like (n_samples, 2048), 
                 'labels': [n_samples, ], 
                 'preds': [n_samples, ], 
                 'subclass_labels': subclass_labels}
    Outputs:
    - results = {'embed': array-like (n_samples, dim), 
                 'labels': [n_samples, ], 
                 'preds': [n_samples, ], 
                 'subclass_labels': subclass_labels}
    """
    with torch.no_grad():
        # svd
        u, s, v = torch.svd(model.fc.weight.T) # 1000, 2048 -> 2048, 1000
        # print(u.shape, s.shape, v.shape)
        s[dim:] = 0
        r = u * s # reduction matrix # 2048, 1000 -> cols past dim are 0
        embed_old = embed_dict['embed'].to(device) # n_samples, 2048
        embed_red = embed_old @ r # n_samples, 1000, w/0s on end
        embed_dict['embed'] = embed_red[:, :dim].cpu()
        embed_dict['preds'] = torch.argmax(embed_red @ v, dim=1).tolist()
        return embed_dict