jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def show_dimensionality_reduction(_dt):
    with st.expander("Dimensionality Reduction"):
        # get head objects.
        W_QK = get_qk_circuit(_dt)
        U_QK, S_QK, V_QK = torch.svd(W_QK)
        W_OV = get_ov_circuit(_dt)
        U_OV, S_OV, V_OV = torch.svd(W_OV)
        W_U = _dt.action_predictor.weight

        a, b = st.columns(2)
        with a:
            selected_writer = st.selectbox(
                "Select Writer",
                options=["Embeddings", "Head Output", "Neuron Output"],
                key="svd_virtual_weights",
            )
        with b:
            st.write("Select the component writing to the residual stream.")

        if selected_writer == "Embeddings":
            keys_tab, queries_tab, values_tab = st.tabs(
                ["Keys", "Queries", "Values"]
            )
            with keys_tab:
                embedding_projection_onto_svd_component(_dt, V_QK, key="keys")

            with queries_tab:
                embedding_projection_onto_svd_component(
                    _dt, U_QK, key="queries"
                )

            with values_tab:
                embedding_projection_onto_svd_component(
                    _dt, V_OV, key="values"
                )

        if selected_writer == "Head Output":
            (
                key_composition_tab,
                query_composition_tab,
                value_composition_tab,
                mlp_in_tab,
                unembedding_tab,
            ) = st.tabs(
                [
                    "Keys",
                    "Queries",
                    "Values",
                    "Neuron Activation",
                    "Unembedding",
                ]
            )

            with key_composition_tab:
                V_Q_tmp = V_QK.permute(0, 1, 3, 2)
                svd_out_to_svd_in_component(
                    _dt, U_OV, V_Q_tmp, key="key composition"
                )

            with query_composition_tab:
                svd_out_to_svd_in_component(
                    _dt, U_OV, U_QK, key="query composition"
                )

            with value_composition_tab:
                V_OV_tmp = V_OV.permute(0, 1, 3, 2)
                svd_out_to_svd_in_component(
                    _dt, U_OV, V_OV_tmp, key="value composition"
                )

            with mlp_in_tab:
                svd_out_to_mlp_in_component(_dt, V_OV)

            with unembedding_tab:
                svd_out_to_unembedding_component(_dt, V_OV, W_U)

        if selected_writer == "Neuron Output":
            key_tab, query_tab, value_tab = st.tabs(
                ["Keys", "Queries", "Values"]
            )

            with key_tab:
                V_QK_tmp = V_QK.permute(0, 1, 3, 2)
                mlp_out_to_svd_in_component(_dt, V_QK_tmp, key="key")

            with query_tab:
                mlp_out_to_svd_in_component(_dt, U_QK, key="query")

            with value_tab:
                V_OV_tmp = V_OV.permute(0, 1, 3, 2)
                mlp_out_to_svd_in_component(_dt, V_OV_tmp, key="value")
//  [python]
def show_dimensionality_reduction(_dt):
    with st.expander("Dimensionality Reduction"):
        # get head objects.
        W_QK = get_qk_circuit(_dt)
        U_QK, S_QK, V_QK = torch.linalg.svd(W_QK)
        W_OV = get_ov_circuit(_dt)
        U_OV, S_OV, V_OV = torch.linalg.svd(W_OV)
        W_U = _dt.action_predictor.weight

        a, b = st.columns(2)
        with a:
            selected_writer = st.selectbox(
                "Select Writer",
                options=["Embeddings", "Head Output", "Neuron Output"],
                key="svd_virtual_weights",
            )
        with b:
            st.write("Select the component writing to the residual stream.")

        if selected_writer == "Embeddings":
            keys_tab, queries_tab, values_tab = st.tabs(
                ["Keys", "Queries", "Values"]
            )
            with keys_tab:
                embedding_projection_onto_svd_component(_dt, V_QK, key="keys")

            with queries_tab:
                embedding_projection_onto_svd_component(
                    _dt, U_QK, key="queries"
                )

            with values_tab:
                embedding_projection_onto_svd_component(
                    _dt, V_OV, key="values"
                )

        if selected_writer == "Head Output":
            (
                key_composition_tab,
                query_composition_tab,
                value_composition_tab,
                mlp_in_tab,
                unembedding_tab,
            ) = st.tabs(
                [
                    "Keys",
                    "Queries",
                    "Values",
                    "Neuron Activation",
                    "Unembedding",
                ]
            )

            with key_composition_tab:
                V_Q_tmp = V_QK.permute(0, 1, 3, 2)
                svd_out_to_svd_in_component(
                    _dt, U_OV, V_Q_tmp, key="key composition"
                )

            with query_composition_tab:
                svd_out_to_svd_in_component(
                    _dt, U_OV, U_QK, key="query composition"
                )

            with value_composition_tab:
                V_OV_tmp = V_OV.permute(0, 1, 3, 2)
                svd_out_to_svd_in_component(
                    _dt, U_OV, V_OV_tmp, key="value composition"
                )

            with mlp_in_tab:
                svd_out_to_mlp_in_component(_dt, V_OV)

            with unembedding_tab:
                svd_out_to_unembedding_component(_dt, V_OV, W_U)

        if selected_writer == "Neuron Output":
            key_tab, query_tab, value_tab = st.tabs(
                ["Keys", "Queries", "Values"]
            )

            with key_tab:
                V_QK_tmp = V_QK.permute(0, 1, 3, 2)
                mlp_out_to_svd_in_component(_dt, V_QK_tmp, key="key")

            with query_tab:
                mlp_out_to_svd_in_component(_dt, U_QK, key="query")

            with value_tab:
                V_OV_tmp = V_OV.permute(0, 1, 3, 2)
                mlp_out_to_svd_in_component(_dt, V_OV_tmp, key="value")
//  [python]
def low_rank_decomposition(weight, reduced_rank=0.15):
    """
    :param          weight: The matrix to decompose, of shape (H, W)
    :param      rank_ratio: rank_of_decomposed_matrix
    :return: L, R
    """

    matrix_dimension = len(weight.size())
    assert matrix_dimension == 2, "Only Support 2D matrix"

    U, S, Vh = torch.svd(weight)
    reduced_rank = int(reduced_rank)

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

    return L, R
//  [python]
def low_rank_decomposition(weight, reduced_rank=0.15):
    """
    :param          weight: The matrix to decompose, of shape (H, W)
    :param      rank_ratio: rank_of_decomposed_matrix
    :return: L, R
    """

    matrix_dimension = len(weight.size())
    assert matrix_dimension == 2, "Only Support 2D matrix"

    U, S, Vh = torch.linalg.svd(weight, full_matrices=False)
    reduced_rank = int(reduced_rank)

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

    return L, R
//  [python]
def test_autograd():

    torch.manual_seed(0)

    n_rep = 5
    d = 4
    x = torch.arange(1, d + 1, dtype=torch.float32)

    W = torch.zeros((d, d)).normal_()
    U, s, V = torch.svd(W)
    W = U

    x.requires_grad_()
    W.requires_grad_()

    with torch.enable_grad():
        y = f_rep(W, x, n_rep)

    d = torch.sum(y)

    grad_output = torch.autograd.grad(outputs=[d], inputs=[y])

    gradient_torch = torch.autograd.grad(
        outputs=[y], inputs=[W, x], grad_outputs=grad_output
    )

    gradient_rev = f_rep_backward(grad_output, W, y, x, n_rep)

    err_0 = abs(gradient_torch[0] - gradient_rev[0]).max()
    err_1 = abs(gradient_torch[1] - gradient_rev[1]).max()
    assert err_0 < 1e-4
    assert err_1 < 1e-4
//  [python]
def test_autograd():

    torch.manual_seed(0)

    n_rep = 5
    d = 4
    x = torch.arange(1, d + 1, dtype=torch.float32)

    W = torch.zeros((d, d)).normal_()
    U, s, V = torch.linalg.svd(W)
    W = U

    x.requires_grad_()
    W.requires_grad_()

    with torch.enable_grad():
        y = f_rep(W, x, n_rep)

    d = torch.sum(y)

    grad_output = torch.autograd.grad(outputs=[d], inputs=[y])

    gradient_torch = torch.autograd.grad(
        outputs=[y], inputs=[W, x], grad_outputs=grad_output
    )

    gradient_rev = f_rep_backward(grad_output, W, y, x, n_rep)

    err_0 = abs(gradient_torch[0] - gradient_rev[0]).max()
    err_1 = abs(gradient_torch[1] - gradient_rev[1]).max()
    assert err_0 < 1e-4
    assert err_1 < 1e-4
//  [python]
def procrustes(X_source, X_target):
    """min |X_source Q - X_target|_F s.t. Q^TQ = I"""
    U, _, Vh = torch.svd(X_target.T @ X_source)
    V = Vh.transpose(-2, -1)
    return V @ U.T
//  [python]
def procrustes(X_source, X_target):
    """min |X_source Q - X_target|_F s.t. Q^TQ = I"""
    U, _, Vh = torch.linalg.svd(X_target.T @ X_source, full_matrices=False)
    V = Vh.transpose(-2, -1)
    return V @ U.T
//  [python]
def _dim_reduce(x, p):
    # PCA to truncate -- project onto top p principal axes (no rescaling)
    with torch.no_grad():
        _, _, vT = torch.svd(x)
    # svd returns v.T, so the principal axes are in the *rows*. The following einsum is equivalent to x @ vT.T[:, :p]
    # but a bit faster because the transpose is not actually performed.
    return torch.einsum("mn,pn->mp", x, vT[:p, :])
//  [python]
def _dim_reduce(x, p):
    # PCA to truncate -- project onto top p principal axes (no rescaling)
    with torch.no_grad():
        _, _, vT = torch.linalg.svd(x, full_matrices=False)
    # svd returns v.T, so the principal axes are in the *rows*. The following einsum is equivalent to x @ vT.T[:, :p]
    # but a bit faster because the transpose is not actually performed.
    return torch.einsum("mn,pn->mp", x, vT[:p, :])
//  [python]
def proj_standardized(X, demean=False, inplace=False):
    if demean:
        if inplace:
            X.sub_(X.mean(axis=0))
        else:
            X -= X.mean(axis=0)

    # pytorch 1.8.0 has a bug in which torch.svd fails when requires_grad
    # is true on the input (even if called under torch.no_grad)
    requires_grad = X.requires_grad
    X.requires_grad_(False)

    n = torch.tensor(X.shape[0], dtype=X.dtype, device=X.device)
    m = X.shape[1]
    # TODO: Gracefully handle the rare svd failure
    # TODO: debug alternative eigenvec approach ...
    # (evals, V = torch.eig(X.T @ X, eigenvectors=True)
    # proj = X @ V @ torch.diag(evals[:, 0].sqrt().pow(-1)) ...
    # proj *= torch.sqrt(n)
    if inplace:
        s = torch.zeros(m, device=X.device, dtype=X.dtype)
        Vh = torch.zeros((m, m), device=X.device, dtype=X.dtype)
        U = torch.zeros(X.shape, device=X.device, dtype=X.dtype)
        try:
            U, s, Vh = torch.svd(X, some=False, out=(U, s, Vh))
            V = Vh.transpose(-2, -1)
        except RuntimeError as e:
            X.requires_grad_(requires_grad)
            raise SolverError(str(e))
        torch.matmul(U[:, :m], V.T[:, :m], out=X)
        X.mul_(torch.sqrt(n))
        X.requires_grad_(requires_grad)
        return X
    else:
        try:
            U, s, Vh = torch.svd(X, some=False)
            V = Vh.transpose(-2, -1)
        except RuntimeError as e:
            X.requires_grad_(requires_grad)
            raise SolverError(str(e))
        proj = torch.sqrt(n) * U[:, :m] @ V.T[:, :m]
        X.requires_grad_(requires_grad)
        return proj
//  [python]
def proj_standardized(X, demean=False, inplace=False):
    if demean:
        if inplace:
            X.sub_(X.mean(axis=0))
        else:
            X -= X.mean(axis=0)

    # pytorch 1.8.0 has a bug in which torch.svd fails when requires_grad
    # is true on the input (even if called under torch.no_grad)
    requires_grad = X.requires_grad
    X.requires_grad_(False)

    n = torch.tensor(X.shape[0], dtype=X.dtype, device=X.device)
    m = X.shape[1]
    # TODO: Gracefully handle the rare svd failure
    # TODO: debug alternative eigenvec approach ...
    # (evals, V = torch.eig(X.T @ X, eigenvectors=True)
    # proj = X @ V @ torch.diag(evals[:, 0].sqrt().pow(-1)) ...
    # proj *= torch.sqrt(n)
    if inplace:
        s = torch.zeros(m, device=X.device, dtype=X.dtype)
        Vh = torch.zeros((m, m), device=X.device, dtype=X.dtype)
        U = torch.zeros(X.shape, device=X.device, dtype=X.dtype)
        try:
            U, _, Vh = torch.linalg.svd(X, full_matrices=False, out=(U, s, Vh))
            V = Vh.transpose(-2, -1)
        except RuntimeError as e:
            X.requires_grad_(requires_grad)
            raise SolverError(str(e))
        torch.matmul(U[:, :m], V.T[:, :m], out=X)
        X.mul_(torch.sqrt(n))
        X.requires_grad_(requires_grad)
        return X
    else:
        try:
            U, _, Vh = torch.linalg.svd(X, full_matrices=False)
            V = Vh.transpose(-2, -1)
        except RuntimeError as e:
            X.requires_grad_(requires_grad)
            raise SolverError(str(e))
        proj = torch.sqrt(n) * U[:, :m] @ V.T[:, :m]
        X.requires_grad_(requires_grad)
        return proj
//  [python]
def _orthogonal_procrustes_rotation(a, b, anchor="middle"):
    """Provided a and b, each matrix of size (m, p) that are already centered and scaled, solve the orthogonal
    procrustest problem (rotate a and b into a common frame that minimizes distances).

    If anchor="middle" (default) then both a and b
    If anchor="a", then a is left unchanged and b is rotated towards it
    If anchor="b", then b is left unchanged and a is rotated towards it

    See https://en.wikipedia.org/wiki/Orthogonal_Procrustes_problem

    :return: r_a and r_b, which, when right-multiplied with a and b, gives the aligned coordinates, or None for each if
    no transform is required
    """
    with torch.no_grad():
        u, _, v = torch.svd(a.T @ b)
    # Helpful trick to see how these are related: u is the inverse of u.T, and likewise v is inverse of v.T. We get to
    # the anchor=a and anchor=b solutions by right-multiplying both return values by u.T or right-multiplying both
    # return values by v, respectively (if both return values are rotated in the same way, it preserves the shape).
    if anchor == "middle":
        return u, v.T
    elif anchor == "a":
        return None, v.T @ u.T
    elif anchor == "b":
        return u @ v, None
    else:
        raise ValueError(f"Invalid 'anchor' argument: {anchor} (must be 'middle', 'a', or 'b')")
//  [python]
def _orthogonal_procrustes_rotation(a, b, anchor="middle"):
    """Provided a and b, each matrix of size (m, p) that are already centered and scaled, solve the orthogonal
    procrustest problem (rotate a and b into a common frame that minimizes distances).

    If anchor="middle" (default) then both a and b
    If anchor="a", then a is left unchanged and b is rotated towards it
    If anchor="b", then b is left unchanged and a is rotated towards it

    See https://en.wikipedia.org/wiki/Orthogonal_Procrustes_problem

    :return: r_a and r_b, which, when right-multiplied with a and b, gives the aligned coordinates, or None for each if
    no transform is required
    """
    with torch.no_grad():
        u, _, v = torch.linalg.svd(a.T @ b)
    # Helpful trick to see how these are related: u is the inverse of u.T, and likewise v is inverse of v.T. We get to
    # the anchor=a and anchor=b solutions by right-multiplying both return values by u.T or right-multiplying both
    # return values by v, respectively (if both return values are rotated in the same way, it preserves the shape).
    if anchor == "middle":
        return u, v.T
    elif anchor == "a":
        return None, v.T @ u.T
    elif anchor == "b":
        return u @ v, None
    else:
        raise ValueError(f"Invalid 'anchor' argument: {anchor} (must be 'middle', 'a', or 'b')")
//  [python]
def kabsch(P: torch.tensor, Q: torch.tensor):
    """Kabsch algorithm to obtain rotation matrix"""
    C = torch.matmul(torch.transpose(P, 0, 1), Q)
    U, S, V = torch.svd(C)
    d = (torch.det(U) * torch.det(V)) < 0.0

    if d:
        S[-1] = -S[-1]
        U[:, -1] = -U[:, -1]

    R = torch.matmul(U, V.T)
    return R
//  [python]
def kabsch(P: torch.tensor, Q: torch.tensor):
    """Kabsch algorithm to obtain rotation matrix"""
    C = torch.matmul(torch.transpose(P, 0, 1), Q)
    V, S, W = torch.linalg.svd(C)
    d = (torch.linalg.det(V) * torch.linalg.det(W)) < 0.0

    if d:
        S[-1] = -S[-1]
        V[:, -1] = -V[:, -1]

    U = torch.matmul(V, W)
    return U
//  [python]
def gen_covariances(
        self,
        ndim: int,
        batch_shape: Sequence[int] = (),
        as_sqrt: bool = False,
    ) -> Tensor:
        shape = tuple(batch_shape) + (ndim, ndim)
        eigvals = -torch.rand(shape[:-1], **self.tkwargs).log()  # exponential rvs
        orthmat = torch.svd(torch.randn(shape, **self.tkwargs)).U
        sqrt_covar = orthmat * torch.sqrt(eigvals).unsqueeze(-2)
        return sqrt_covar if as_sqrt else sqrt_covar @ sqrt_covar.transpose(-2, -1)
//  [python]
def gen_covariances(
        self,
        ndim: int,
        batch_shape: Sequence[int] = (),
        as_sqrt: bool = False,
    ) -> Tensor:
        shape = tuple(batch_shape) + (ndim, ndim)
        eigvals = -torch.rand(shape[:-1], **self.tkwargs).log()  # exponential rvs
        orthmat = torch.linalg.svd(torch.randn(shape, **self.tkwargs)).U
        sqrt_covar = orthmat * torch.sqrt(eigvals).unsqueeze(-2)
        return sqrt_covar if as_sqrt else sqrt_covar @ sqrt_covar.transpose(-2, -1)
//  [python]
def neural_data_to_point(self, x: NeuralData) -> Point:
        """Convert size (m,d) neural data to a size (m,m) Gram matrix of inner products between xs using self._kernel.
        """
        # Always preprocess by subtracting the mean to ensure translation invariance
        x = x - x.mean(dim=0, keepdims=True)
        if x.shape[0] != self.m:
            raise ValueError(f"Expected x to be size ({self.m}, ?) but is size {x.shape}")
        if self._mode == "gram":
            gram_matrix = pairwise.inner_product(x, kernel=self._kernel)
            # Apply regularizer: add a small ridge to the diagonal
            return gram_matrix + self._eps * torch.eye(self.m, device=x.device, dtype=x.dtype)
        elif self._mode == "cov":
            # Flatten all but first dimension.
            x = torch.reshape(x, (self.m, -1))
            # Pad or truncate to p dimensions
            d = prod(x.shape) // self.m
            if d > self._p:
                # PCA to truncate -- project onto top p principal axes (no rescaling)
                u, s, v = torch.svd(x)
                x = x @ v[:, :self._p]
            elif d < self._p:
                # Pad zeros
                num_pad = self._p - d
                x = torch.hstack([x.view(self.m, d), x.new_zeros(self.m, num_pad)])
            # Compute covariance
            cov_matrix = torch.einsum('mi,mj->ij', x, x) / (self.m - 1)
            # Apply regularizer: add small ridge
            return cov_matrix + self._eps * torch.eye(self._p, device=x.device, dtype=x.dtype)
//  [python]
def neural_data_to_point(self, x: NeuralData) -> Point:
        """Convert size (m,d) neural data to a size (m,m) Gram matrix of inner products between xs using self._kernel.
        """
        # Always preprocess by subtracting the mean to ensure translation invariance
        x = x - x.mean(dim=0, keepdims=True)
        if x.shape[0] != self.m:
            raise ValueError(f"Expected x to be size ({self.m}, ?) but is size {x.shape}")
        if self._mode == "gram":
            gram_matrix = pairwise.inner_product(x, kernel=self._kernel)
            # Apply regularizer: add a small ridge to the diagonal
            return gram_matrix + self._eps * torch.eye(self.m, device=x.device, dtype=x.dtype)
        elif self._mode == "cov":
            # Flatten all but first dimension.
            x = torch.reshape(x, (self.m, -1))
            # Pad or truncate to p dimensions
            d = prod(x.shape) // self.m
            if d > self._p:
                # PCA to truncate -- project onto top p principal axes (no rescaling)
                _, _, v = svd(x)
                x = x @ v[:, :self._p]
            elif d < self._p:
                # Pad zeros
                num_pad = self._p - d
                x = torch.hstack([x.view(self.m, d), x.new_zeros(self.m, num_pad)])
            # Compute covariance
            cov_matrix = torch.einsum('mi,mj->ij', x, x) / (self.m - 1)
            # Apply regularizer: add small ridge
            return cov_matrix + self._eps * torch.eye(self._p, device=x.device, dtype=x.dtype)
//  [python]
def solve_least_squares(A, B):
    # print(A.shape)
    # print(B.shape)
    # Compute C = A^T A
    # min_eigenvalues = torch.min( torch.linalg.eigvalsh(C), dim=-1 )
    # eps_e = torch.maximum( min_eigenvalues, min_eigenvalues.new_ones(min_eigenvalues.shape)*1e-3 )[...,]
    C = torch.matmul(A.transpose(-2, -1), A)  # + eps_e*torch.eye(A.shape[-1], device=A.device)
    # Compute the pseudo-inverse of C
    U, S, Vh = torch.svd(C.float())
    D_inv = torch.diag_embed(1.0 / torch.maximum(S, torch.ones_like(S) * 1e-4))
    C_inv = Vh.transpose(-1,-2).matmul(D_inv).matmul(U.transpose(-1,-2))

    # Compute X = C_inv A^T B
    X = torch.matmul(torch.matmul(C_inv, A.transpose(-2, -1)), B)
    return X
//  [python]
def solve_least_squares(A, B):
    # print(A.shape)
    # print(B.shape)
    # Compute C = A^T A
    # min_eigenvalues = torch.min( torch.linalg.eigvalsh(C), dim=-1 )
    # eps_e = torch.maximum( min_eigenvalues, min_eigenvalues.new_ones(min_eigenvalues.shape)*1e-3 )[...,]
    C = torch.matmul(A.transpose(-2, -1), A)  # + eps_e*torch.eye(A.shape[-1], device=A.device)
    # Compute the pseudo-inverse of C
    U, S, Vh = torch.linalg.svd(C.float(), full_matrices=False)
    D_inv = torch.diag_embed(1.0 / torch.maximum(S, torch.ones_like(S) * 1e-4))
    C_inv = Vh.transpose(-1,-2).matmul(D_inv).matmul(U.transpose(-1,-2))

    # Compute X = C_inv A^T B
    X = torch.matmul(torch.matmul(C_inv, A.transpose(-2, -1)), B)
    return X
//  [python]
def extract_linear(weight, lora_rank, dynamic_method, dynamic_param, device, scale=1):
    out_size, in_size = weight.size()

    U, S, Vh = torch.svd(weight.to(device))

    param_dict = rank_resize(S, lora_rank, dynamic_method, dynamic_param, scale)
    lora_rank = param_dict["new_rank"]

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

    param_dict["lora_down"] = Vh.reshape(lora_rank, in_size).cpu()
    param_dict["lora_up"] = U.reshape(out_size, lora_rank).cpu()
    del U, S, Vh, weight
    return param_dict
//  [python]
def extract_linear(weight, lora_rank, dynamic_method, dynamic_param, device, scale=1):
    out_size, in_size = weight.size()

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

    param_dict = rank_resize(S, lora_rank, dynamic_method, dynamic_param, scale)
    lora_rank = param_dict["new_rank"]

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

    param_dict["lora_down"] = Vh.reshape(lora_rank, in_size).cpu()
    param_dict["lora_up"] = U.reshape(out_size, lora_rank).cpu()
    del U, S, Vh, weight
    return param_dict
//  [python]
def RR(self, rangles, observations):
        """
        analytically solve the weights/coefficients from observations using RR
        """
        eta = 40
        th, td, pd = torch.unbind(rangles, dim=1)
        stacks = PCA.unmask(
            torch.hstack([self.Q, self.median, self.relativeOffset]), self.maskMap
        )

        Q, median, relativeOffset = torch.split(
            Merl.merl_lookup(stacks.T, th, td, pd, scaling=False).T,
            [self.n, 1, 1],
            dim=1,
        )
        ph = torch.zeros_like(th).to(th.device)
        wi, wo = coords.hd_to_io_sph(
            torch.stack([th, ph], dim=1), torch.stack([td, pd], dim=1)
        )
        cosMap = torch.cos(wi[:, [0]]) * torch.cos(wo[:, [0]])
        cosMap[cosMap < 0.0] = 0.0  # max(cos, 0.0)

        mappedObs = torch.log((observations * cosMap + 1e-4) / (median + 1e-4))
        b = mappedObs - relativeOffset
        U, s, Vt = torch.svd(Q)
        sinv = torch.diag(s / (s * s + eta))

        c = Vt.T @ sinv @ U.T @ b

        return c
//  [python]
def RR(self, rangles, observations):
        """
        analytically solve the weights/coefficients from observations using RR
        """
        eta = 40
        th, td, pd = torch.unbind(rangles, dim=1)
        stacks = PCA.unmask(
            torch.hstack([self.Q, self.median, self.relativeOffset]), self.maskMap
        )

        Q, median, relativeOffset = torch.split(
            Merl.merl_lookup(stacks.T, th, td, pd, scaling=False).T,
            [self.n, 1, 1],
            dim=1,
        )
        ph = torch.zeros_like(th).to(th.device)
        wi, wo = coords.hd_to_io_sph(
            torch.stack([th, ph], dim=1), torch.stack([td, pd], dim=1)
        )
        cosMap = torch.cos(wi[:, [0]]) * torch.cos(wo[:, [0]])
        cosMap[cosMap < 0.0] = 0.0  # max(cos, 0.0)

        mappedObs = torch.log((observations * cosMap + 1e-4) / (median + 1e-4))
        b = mappedObs - relativeOffset
        U, s, Vt = torch.linalg.svd(Q, full_matrices=False)
        sinv = torch.diag(s / (s * s + eta))

        c = Vt.T @ sinv @ U.T @ b

        return c
//  [python]
def generate_orthogonal_embs(total_spks, perturb_sigma, emb_dim):
    """Generate a set of artificial orthogonal embedding vectors from random numbers
    """
    gaus = torch.randn(emb_dim, emb_dim)
    _svd = torch.svd(gaus)
    orth = _svd[0] @ _svd[2]
    orth_embs = orth[:total_spks]
    # Assert orthogonality
    assert torch.abs(getCosAffinityMatrix(orth_embs) - torch.diag(torch.ones(total_spks))).sum() < 1e-4
    return orth_embs
//  [python]
def generate_orthogonal_embs(total_spks, perturb_sigma, emb_dim):
    """Generate a set of artificial orthogonal embedding vectors from random numbers
    """
    gaus = torch.randn(emb_dim, emb_dim)
    _svd = torch.linalg.svd(gaus)
    orth = _svd[0] @ _svd[2]
    orth_embs = orth[:total_spks]
    # Assert orthogonality
    assert torch.abs(getCosAffinityMatrix(orth_embs) - torch.diag(torch.ones(total_spks))).sum() < 1e-4
    return orth_embs
//  [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 f"model.layers.{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.float()
                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
                reconstructed_weights = reconstructed_weights.to(original_dtype)
                module.weight = torch.nn.Parameter(reconstructed_weights)
                self.modified_layers.add(layer_id)

                del U, S, V, weights
                torch.mps.empty_cache()  # Clear PyTorch's MPS memory cache
                gc.collect()
                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 f"model.layers.{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.float()
                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)

                del U, S, V, weights
                torch.mps.empty_cache()  # Clear PyTorch's MPS memory cache
                gc.collect()
                return True
//  [python]
def _low_rank_decomposition(
    weight: torch.Tensor, reduced_rank: int = 16
) -> Tuple[torch.Tensor, torch.Tensor]:
    """
    Decompose a 2D matrix into low-rank matrices A and B using SVD.a

    :param weight: The matrix to decompose, of shape (H, W)
    :param reduced_rank: The final rank of the decomposition
    :return: A tuple of tensors (A, B)
    """
    if weight.dim() != 2:
        raise ValueError(
            f"Only support 2D matrix, but your input has {weight.dim()} dimensions."
        )

    # SVD Decomposition
    U, S, Vh = torch.svd(weight)

    # Truncated matrices
    A = Vh[:reduced_rank, :]
    B = U[:, :reduced_rank] @ torch.diag(S[:reduced_rank])

    return A, B
//  [python]
def _low_rank_decomposition(
    weight: torch.Tensor, reduced_rank: int = 16
) -> Tuple[torch.Tensor, torch.Tensor]:
    """
    Decompose a 2D matrix into low-rank matrices A and B using SVD.a

    :param weight: The matrix to decompose, of shape (H, W)
    :param reduced_rank: The final rank of the decomposition
    :return: A tuple of tensors (A, B)
    """
    if weight.dim() != 2:
        raise ValueError(
            f"Only support 2D matrix, but your input has {weight.dim()} dimensions."
        )

    # SVD Decomposition
    U, S, Vh = torch.linalg.svd(weight, full_matrices=False)

    # Truncated matrices
    A = Vh[:reduced_rank, :]
    B = U[:, :reduced_rank] @ torch.diag(S[:reduced_rank])

    return A, B
//  [python]
def svd_orthogonal(fin, fout, use_paddle=False):
    assert fin > fout, 'fin > fout'
    if use_paddle:
        X = torch.rand((fout, fin))
        U, S, Vt = torch.svd(X)
        # print(Vt.shape)
        # print(paddle.allclose(Vt@Vt.T, paddle.eye(Vt.shape[0])))
    else:
        X = np.random.random((fout, fin))
        U, S, Vt = np.linalg.svd(X, full_matrices=False)
        # print(Vt.shape)
        # print(np.allclose((Vt@ Vt.T), np.eye(Vt.shape[0])))
    W = torch.tensor(Vt, dtype=torch.float32).T
    return W
//  [python]
def svd_orthogonal(fin, fout, use_paddle=False):
    assert fin > fout, 'fin > fout'
    if use_paddle:
        X = torch.rand((fout, fin))
        U, _, Vt = torch.linalg.svd(X, full_matrices=False)
        # print(Vt.shape)
        # print(paddle.allclose(Vt@Vt.T, paddle.eye(Vt.shape[0])))
    else:
        X = np.random.random((fout, fin))
        U, _, Vt = np.linalg.svd(X, full_matrices=False)
        # print(Vt.shape)
        # print(np.allclose((Vt@ Vt.T), np.eye(Vt.shape[0])))
    W = torch.tensor(Vt, dtype=torch.float32).T
    return W
//  [python]
def normalizeSE3_inplace(T):
  R = T[...,:3,:3]
  U, S, Vh = torch.svd(R)
  T[...,:3,:3] = torch.matmul(U, Vh)
//  [python]
def normalizeSE3_inplace(T):
  R = T[...,:3,:3]
  U, S, Vh = torch.linalg.svd(R)
  T[...,:3,:3] = torch.matmul(U, Vh)
//  [python]
def pca_forward(self, x):
        batch_size, feature_dim = x.size()
        f_cov = (torch.mm(x.transpose(0, 1), x) / (batch_size - 1)).float()  # d * d
        eye = torch.eye(feature_dim).float().to(f_cov.device)
        # f_cov = torch.mm(x.transpose(0, 1), x).float()  # d * d
        U, S, V = torch.svd(0.99 * f_cov + 0.01 * eye)
        diag = torch.diag(1.0 / torch.sqrt(S + 1e-5))
        rotate_mtx = torch.mm(diag, U.transpose(0, 1)).detach()  # d * d
        return torch.mm(rotate_mtx, x.transpose(0, 1)).transpose(0, 1)
//  [python]
def pca_forward(self, x):
        batch_size, feature_dim = x.size()
        f_cov = (torch.mm(x.transpose(0, 1), x) / (batch_size - 1)).float()  # d * d
        eye = torch.eye(feature_dim).float().to(f_cov.device)
        # f_cov = torch.mm(x.transpose(0, 1), x).float()  # d * d
        U, S, V = torch.linalg.svd(0.99 * f_cov + 0.01 * eye)
        diag = torch.diag(1.0 / torch.sqrt(S + 1e-5))
        rotate_mtx = torch.mm(diag, U.transpose(0, 1)).detach()  # d * d
        return torch.mm(rotate_mtx, x.transpose(0, 1)).transpose(0, 1)
//  [python]
def singular_value_operation(
    number, transform_coeff, operation, cast_to_complex=True, wandb_name=None
):
    U, S, V = torch.svd(transform_coeff)
    D = torch.diag(S)
    D_clipped = operation(D, number)
    if wandb_name is not None:
        wandb.log({wandb_name: wandb.Histogram(D)})
    if cast_to_complex:
        D_clipped = D_clipped.type(torch.complex64)
    D_clipped = torch.diag_embed(D_clipped)
    clipped_coeff = torch.matmul(D_clipped, V)
    clipped_coeff = torch.matmul(U, clipped_coeff)
    return clipped_coeff
//  [python]
def singular_value_operation(
    number, transform_coeff, operation, cast_to_complex=True, wandb_name=None
):
    U, D, V = torch.linalg.svd(transform_coeff, full_matrices=False)
    D_clipped = operation(D, number)
    if wandb_name is not None:
        wandb.log({wandb_name: wandb.Histogram(D)})
    if cast_to_complex:
        D_clipped = D_clipped.type(torch.complex64)
    D_clipped = torch.diag_embed(D_clipped)
    clipped_coeff = torch.linalg.matmul(D_clipped, V)
    clipped_coeff = torch.linalg.matmul(U, clipped_coeff)
    return clipped_coeff
//  [python]
def orthogonal_gradient(self):
        weight_grad = self.netG.layers[0].layers[0].weight.grad
        B, C, H, W = weight_grad.size()
        u, s, vh = torch.svd(weight_grad.view(B, C * H * W))
        s[:] = 1
        self.netG.layers[0].layers[0].weight.grad = u.mm(torch.diag(s)).mm(vh).view(B, C, H, W)
//  [python]
def orthogonal_gradient(self):
        weight_grad = self.netG.layers[0].layers[0].weight.grad
        B, C, H, W = weight_grad.size()
        u, s, vh = torch.linalg.svd(weight_grad.view(B, C * H * W),full_matrices=False)
        s[:] = 1
        self.netG.layers[0].layers[0].weight.grad = u.mm(torch.diag(s)).mm(vh).view(B, C, H, W)
//  [python]
def SVD(mat):
    """
    Computes the SVD of a matrix

    Args:
        mat (torch.tensor): the matrix

    Returns:
        U, S, V: the SVD factors.
    """
    if mat.shape[0] < 10*mat.shape[1]:
        try:
            u, s, v = t.svd(mat,full_matrices=False)
            s = s.to(v.dtype)
            return u, s, v
        except:
            u, s, v = np.linalg.svd(mat.numpy(),full_matrices=False)
            return tn.tensor(u, dtype = mat.dtype, device = mat.device), tn.tensor(s, dtype = mat.dtype, device = mat.device), tn.tensor(v, dtype = mat.dtype, device = mat.device)
    else:
        try:    
            u, s, v = t.svd(mat.t(),full_matrices=False)
            s = s.to(v.dtype)
            return  v.t(), s, u.t()
        except:
            u, s, v = np.linalg.svd((mat.t()).numpy(),full_matrices=False)
            return  tn.tensor(v.t(), dtype = mat.dtype, device = mat.device), tn.tensor(s, dtype = mat.dtype, device = mat.device), tn.tensor(u.t(), dtype = mat.dtype, device = mat.device)
//  [python]
def SVD(mat):
    """
    Computes the SVD of a matrix

    Args:
        mat (torch.tensor): the matrix

    Returns:
        U, S, V: the SVD factors.
    """
    if mat.shape[0] < 10*mat.shape[1]:
        try:
            u, s, v = tn.linalg.svd(mat,full_matrices=False)
            s = s.to(v.dtype)
            return u, s, v
        except:
            u, s, v = np.linalg.svd(mat.numpy(),full_matrices=False)
            return tn.tensor(u, dtype = mat.dtype, device = mat.device), tn.tensor(s, dtype = mat.dtype, device = mat.device), tn.tensor(v, dtype = mat.dtype, device = mat.device)
    else:
        try:    
            u, s, v = tn.linalg.svd(mat.t(),full_matrices=False)
            s = s.to(v.dtype)
            return  v.t(), s, u.t()
        except:
            u, s, v = np.linalg.svd((mat.t()).numpy(),full_matrices=False)
            return  tn.tensor(v.t(), dtype = mat.dtype, device = mat.device), tn.tensor(s, dtype = mat.dtype, device = mat.device), tn.tensor(u.t(), dtype = mat.dtype, device = mat.device)
//  [python]
def rot9d_to_rotmat(x):
    """Convert 9D rotation representation to 3x3 rotation matrix.
    Based on Levinson et al., "An Analysis of SVD for Deep Rotation Estimation"
    Input:
        (B,9) or (B,J*9) Batch of 9D rotation (interpreted as 3x3 est rotmat)
    Output:
        (B,3,3) or (B*J,3,3) Batch of corresponding rotation matrices
    """
    x = x.view(-1,3,3)
    u, s, vh = torch.svd(x)

    sig = torch.eye(3).expand(len(x), 3, 3).clone()
    sig = sig.to(x.device)
    sig[:, -1, -1] = (u @ vh).det()

    R = u @ sig @ vh

    return R
//  [python]
def rot9d_to_rotmat(x):
    """Convert 9D rotation representation to 3x3 rotation matrix.
    Based on Levinson et al., "An Analysis of SVD for Deep Rotation Estimation"
    Input:
        (B,9) or (B,J*9) Batch of 9D rotation (interpreted as 3x3 est rotmat)
    Output:
        (B,3,3) or (B*J,3,3) Batch of corresponding rotation matrices
    """
    x = x.view(-1,3,3)
    u, _, vh = torch.linalg.svd(x)

    sig = torch.eye(3).expand(len(x), 3, 3).clone()
    sig = sig.to(x.device)
    sig[:, -1, -1] = (u @ vh).det()

    R = u @ sig @ vh

    return R
//  [python]
def Tchernychova_Lyons_CAR(X, mu, tm, DEBUG=False):
    """
    This functions reduce X from N points to n+1.
    This is taken from https://github.com/FraCose/Recombination_Random_Algos/blob/master/recombination.py
    """
    X = torch.cat([tm.ones(X.size(0)).unsqueeze(0).T, X], dim=1)
    N, n = X.shape
    U, Sigma, V = torch.svd(X.T)
    U = torch.cat([U, tm.zeros(n, N - n)], dim=1)
    Sigma = torch.cat([Sigma, tm.zeros(N - n)])
    Phi = V[-(N - n):, :].T
    cancelled = tm.null()

    for _ in range(N - n):
        lm = len(mu)
        plis = Phi[:, 0] > 0
        # Added 7th Aug, 2023.
        if plis.sum() == 0:
            break
        
        alpha = tm.zeros(lm)
        alpha[plis] = mu[plis] / Phi[plis, 0]
        idx = tm.arange(lm)[plis]
        idx = idx[torch.argmin(alpha[plis])]

        if len(cancelled) == 0:
            cancelled = idx.unsqueeze(0)
        else:
            cancelled = torch.cat([cancelled, idx.unsqueeze(0)])
        mu[:] = mu - alpha[idx] * Phi[:, 0]
        mu[idx] = 0.

        if DEBUG and (not torch.allclose(torch.sum(mu), 1.)):
            # print("ERROR")
            print("sum ", torch.sum(mu))

        Phi_tmp = Phi[:, 0]
        Phi = Phi[:, 1:]
        Phi = Phi - torch.matmul(
            Phi[idx].unsqueeze(1),
            Phi_tmp.unsqueeze(1).T,
        ).T / Phi_tmp[idx]
        Phi[idx, :] = 0.

    w_star = mu[mu > 0]
    idx_star = tm.arange(N)[mu > 0]
    return w_star, idx_star, torch.nan, torch.nan, 0., torch.nan, torch.nan
//  [python]
def Tchernychova_Lyons_CAR(X, mu, tm, DEBUG=False):
    """
    This functions reduce X from N points to n+1.
    This is taken from https://github.com/FraCose/Recombination_Random_Algos/blob/master/recombination.py
    """
    X = torch.cat([tm.ones(X.size(0)).unsqueeze(0).T, X], dim=1)
    N, n = X.shape
    U, Sigma, V = torch.linalg.svd(X.T)
    U = torch.cat([U, tm.zeros(n, N - n)], dim=1)
    Sigma = torch.cat([Sigma, tm.zeros(N - n)])
    Phi = V[-(N - n):, :].T
    cancelled = tm.null()

    for _ in range(N - n):
        lm = len(mu)
        plis = Phi[:, 0] > 0
        # Added 7th Aug, 2023.
        if plis.sum() == 0:
            break
        
        alpha = tm.zeros(lm)
        alpha[plis] = mu[plis] / Phi[plis, 0]
        idx = tm.arange(lm)[plis]
        idx = idx[torch.argmin(alpha[plis])]

        if len(cancelled) == 0:
            cancelled = idx.unsqueeze(0)
        else:
            cancelled = torch.cat([cancelled, idx.unsqueeze(0)])
        mu[:] = mu - alpha[idx] * Phi[:, 0]
        mu[idx] = 0.

        if DEBUG and (not torch.allclose(torch.sum(mu), 1.)):
            # print("ERROR")
            print("sum ", torch.sum(mu))

        Phi_tmp = Phi[:, 0]
        Phi = Phi[:, 1:]
        Phi = Phi - torch.matmul(
            Phi[idx].unsqueeze(1),
            Phi_tmp.unsqueeze(1).T,
        ).T / Phi_tmp[idx]
        Phi[idx, :] = 0.

    w_star = mu[mu > 0]
    idx_star = tm.arange(N)[mu > 0]
    return w_star, idx_star, torch.nan, torch.nan, 0., torch.nan, torch.nan
//  [python]
def __init__(self, num_embeddings, embedding_dim, padding_idx=None):
        super(SparseEmbedding, self).__init__(num_embeddings, embedding_dim, padding_idx=padding_idx)
        self.embedding_dim_origin = embedding_dim
        self.embedding_dim_sparsified = embedding_dim

        u, s, vh = torch.svd(self.weight.data.clone().detach())
        self.u = nn.Parameter(u) 
        self.s = nn.Parameter(s) 
        self.vh = nn.Parameter(vh)
//  [python]
def __init__(self, num_embeddings, embedding_dim, padding_idx=None):
        super(SparseEmbedding, self).__init__(num_embeddings, embedding_dim, padding_idx=padding_idx)
        self.embedding_dim_origin = embedding_dim
        self.embedding_dim_sparsified = embedding_dim

        u, s, vh = torch.linalg.svd(self.weight.data.clone().detach(), full_matrices=False)
        self.u = nn.Parameter(u) 
        self.s = nn.Parameter(s) 
        self.vh = nn.Parameter(vh)
//  [python]
def query_database(scale, z_so3, z_inv, database, chunk=32):
    cur = 0
    knn_pcl_aligned, optimal_R = [], []
    knn_scale, nn_dist, nn_info_bbox = [], [], []
    while cur < len(scale):
        _scale = scale[cur : cur + chunk]
        _z_so3, _z_inv = z_so3[cur : cur + chunk], z_inv[cur : cur + chunk]

        # ! WARNING: THE SOLVED BASIS/R IS IN O(3) NOT ALWAYS SO(3) because it's possible det(R) = -1
        device = _z_so3.device
        # * Query NN from database
        nn_idx, nn_info, _nn_dist = database.query(
            z_so3=_z_so3, z_inv=_z_inv, device=device, use_inv=True, use_so3=True, use_joint=True
        )  # B,K
        # * Align NN PCL to pred PCL with latent registration
        knn_pcl_centered = nn_info["pcl"] - nn_info["center"]  # B,K,N,3
        knn_z_so3 = nn_info["z_so3"]  # B,K,C,3; z_so3 [B,C,3]
        _knn_scale = nn_info["scale"]
        latent_W = (_z_so3[:, None, :, :, None] @ knn_z_so3.unsqueeze(-2)).sum(2)
        u, s, vh = torch.svd(latent_W.double())
        _optimal_R = (u @ vh).float()  # B,K,3,3  R @ knn -> obs
        knn_pcl_rotated = _optimal_R[:, :, None, :, :] @ knn_pcl_centered[..., None]
        knn_pcl_rotated = knn_pcl_rotated.squeeze(-1)  # B,K,N,3
        _knn_pcl_aligned = (
            knn_pcl_rotated / _knn_scale[..., None, None] * _scale[:, None, None, None]
        )

        knn_pcl_aligned.append(_knn_pcl_aligned)
        optimal_R.append(_optimal_R)
        knn_scale.append(_knn_scale)
        nn_dist.append(_nn_dist)
        nn_info_bbox.append(nn_info["bbox"])
        cur += chunk

    knn_pcl_aligned = torch.cat(knn_pcl_aligned, 0)
    optimal_R = torch.cat(optimal_R, 0)
    knn_scale = torch.cat(knn_scale, 0)
    nn_dist = np.concatenate(nn_dist, 0)
    nn_info_bbox = torch.cat(nn_info_bbox, 0)

    return knn_pcl_aligned, optimal_R, knn_scale, nn_dist, nn_info_bbox
//  [python]
def query_database(scale, z_so3, z_inv, database, chunk=32):
    cur = 0
    knn_pcl_aligned, optimal_R = [], []
    knn_scale, nn_dist, nn_info_bbox = [], [], []
    while cur < len(scale):
        _scale = scale[cur : cur + chunk]
        _z_so3, _z_inv = z_so3[cur : cur + chunk], z_inv[cur : cur + chunk]

        # ! WARNING: THE SOLVED BASIS/R IS IN O(3) NOT ALWAYS SO(3) because it's possible det(R) = -1
        device = _z_so3.device
        # * Query NN from database
        nn_idx, nn_info, _nn_dist = database.query(
            z_so3=_z_so3, z_inv=_z_inv, device=device, use_inv=True, use_so3=True, use_joint=True
        )  # B,K
        # * Align NN PCL to pred PCL with latent registration
        knn_pcl_centered = nn_info["pcl"] - nn_info["center"]  # B,K,N,3
        knn_z_so3 = nn_info["z_so3"]  # B,K,C,3; z_so3 [B,C,3]
        _knn_scale = nn_info["scale"]
        latent_W = (_z_so3[:, None, :, :, None] @ knn_z_so3.unsqueeze(-2)).sum(2)
        u, _, vh = torch.linalg.svd(latent_W.double())
        _optimal_R = (u @ vh).float()  # B,K,3,3  R @ knn -> obs
        knn_pcl_rotated = _optimal_R[:, :, None, :, :] @ knn_pcl_centered[..., None]
        knn_pcl_rotated = knn_pcl_rotated.squeeze(-1)  # B,K,N,3
        _knn_pcl_aligned = (
            knn_pcl_rotated / _knn_scale[..., None, None] * _scale[:, None, None, None]
        )

        knn_pcl_aligned.append(_knn_pcl_aligned)
        optimal_R.append(_optimal_R)
        knn_scale.append(_knn_scale)
        nn_dist.append(_nn_dist)
        nn_info_bbox.append(nn_info["bbox"])
        cur += chunk

    knn_pcl_aligned = torch.cat(knn_pcl_aligned, 0)
    optimal_R = torch.cat(optimal_R, 0)
    knn_scale = torch.cat(knn_scale, 0)
    nn_dist = np.concatenate(nn_dist, 0)
    nn_info_bbox = torch.cat(nn_info_bbox, 0)

    return knn_pcl_aligned, optimal_R, knn_scale, nn_dist, nn_info_bbox
//  [python]
def svd_decomposition(
    A: torch.Tensor,
    randomized: bool,
    num_ranks: int,
    num_oversampling: int = 5,
) -> Tuple[torch.Tensor, torch.Tensor]:
    if A.ndim != 2:
        raise ValueError(f"Expected 2D Matrix, but got {A.ndim}.")

    if randomized is False:
        U, S, VT = torch.svd(A, some=False)
    elif randomized is True:
        U, S, V = torch.svd_lowrank(A, num_ranks + num_oversampling)
        # https://pytorch.org/docs/stable/_modules/torch/_lowrank.html#svd_lowrank
        VT = V.mH
    else:
        raise ValueError(f"`randomized` {randomized} not supported")

    S_sqrt = torch.sqrt(S)
    L1 = U * S_sqrt.unsqueeze(dim=0)
    L2 = VT * S_sqrt.unsqueeze(dim=1)
    L1k = L1[:, :num_ranks]
    L2k = L2[:num_ranks, :]
    return L1k, L2k
//  [python]
def svd_decomposition(
    A: torch.Tensor,
    randomized: bool,
    num_ranks: int,
    num_oversampling: int = 5,
) -> Tuple[torch.Tensor, torch.Tensor]:
    if A.ndim != 2:
        raise ValueError(f"Expected 2D Matrix, but got {A.ndim}.")

    if randomized is False:
        U, S, VT = torch.linalg.svd(A, full_matrices=False)
    elif randomized is True:
        U, S, V = torch.svd_lowrank(A, num_ranks + num_oversampling)
        # https://pytorch.org/docs/stable/_modules/torch/_lowrank.html#svd_lowrank
        VT = V.mH
    else:
        raise ValueError(f"`randomized` {randomized} not supported")

    S_sqrt = torch.sqrt(S)
    L1 = U * S_sqrt.unsqueeze(dim=0)
    L2 = VT * S_sqrt.unsqueeze(dim=1)
    L1k = L1[:, :num_ranks]
    L2k = L2[:num_ranks, :]
    return L1k, L2k