jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def __init__(self, *,
                 num_classes: int,
                 hidden_dim: int = 16,
                 head_layers: int = 1,
                 dropout: float = 0.0,
                 activation_fn: Callable[[Tensor], Tensor] = torch.relu_,
                 batch_norm: bool = False,
                 norm_scale: float = 1.0,
                 ):

        super().__init__(
            num_classes=num_classes,
            hidden_dim=hidden_dim,
            num_layers=head_layers,
            dropout=dropout,
            activation_fn=activation_fn,
            batch_norm=batch_norm,
        )
        
        # The MLP_A is hard-coded to be a one layer MLP (XW^T+b)
        self.mlp_adj = Linear(-1, hidden_dim)
        self.norm_scale = norm_scale
        
        # Freeze the head MLP as initialization, and disable bias
        for lin in self.model.layers:
            lin.bias = torch.nn.init.constant_(lin.bias, 0.0)
            # we use random orthonormal matrix to construct projection. It prevents potential rank collapse.
            Rot = torch.randn(hidden_dim,hidden_dim)
            svd = torch.linalg.svd(Rot)
            orth = svd[0] @ svd[2]
            lin.weight = Parameter(torch.eye(num_classes,hidden_dim) @ orth * math.sqrt(3.0/hidden_dim*2))
        for param in self.model.parameters():
            param.requires_grad = False
        
        # extract the weight matrix and do column normalization of W (row normalization of W^T)
        # self.mlp_adj.state_dict()['weight'] = F.normalize(self.mlp_adj.state_dict()['weight'], p=2, dim=0)
        # debug use only
        # print(torch.norm(self.mlp_adj.state_dict()['weight'], dim=0))
        
        self.dropout_fn = Dropout(p=dropout, inplace=True)
        self.activation_fn = activation_fn
        self.bn = BatchNorm1d(hidden_dim) if batch_norm else False
//  [python]
def Tchernychova_Lyons_CAR(X, mu, device, 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([torch.ones(X.size(0)).unsqueeze(0).T.to(device), X], dim=1)
    N, n = X.shape
    U, Sigma, V = torch.svd(X.T)
    U = torch.cat([U, torch.zeros((n, N - n)).to(device)], dim=1)
    Sigma = torch.cat([Sigma, torch.zeros(N - n).to(device)])
    Phi = V[-(N - n):, :].T
    cancelled = torch.tensor([], dtype=int).to(device)

    for _ in range(N - n):
        lm = len(mu)
        plis = Phi[:, 0] > 0
        alpha = torch.zeros(lm).to(device)
        alpha[plis] = mu[plis] / Phi[plis, 0]
        idx = torch.arange(lm)[plis].to(device)
        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 = torch.arange(N)[mu > 0].to(device)
    return w_star, idx_star, torch.nan, torch.nan, 0., torch.nan, torch.nan
//  [python]
def Tchernychova_Lyons_CAR(X, mu, device, 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([torch.ones(X.size(0)).unsqueeze(0).T.to(device), X], dim=1)
    N, n = X.shape
    U, Sigma, V = torch.linalg.svd(X.T)
    U = torch.cat([U, torch.zeros((n, N - n)).to(device)], dim=1)
    Sigma = torch.cat([Sigma, torch.zeros(N - n).to(device)])
    Phi = V[-(N - n):, :].T
    cancelled = torch.tensor([], dtype=int).to(device)

    for _ in range(N - n):
        lm = len(mu)
        plis = Phi[:, 0] > 0
        alpha = torch.zeros(lm).to(device)
        alpha[plis] = mu[plis] / Phi[plis, 0]
        idx = torch.arange(lm)[plis].to(device)
        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 = torch.arange(N)[mu > 0].to(device)
    return w_star, idx_star, torch.nan, torch.nan, 0., torch.nan, torch.nan
//  [python]
def forward(self, X: Tensor, center_data: bool) -> Tuple[Tensor, Tensor]:
        """
        Perform PCA on the data matrix X by computing its SVD.

        Args:
            X (Tensor): Data matrix.
            center_data (bool): If true, then the data mean will be subtracted
            from all data points prior to performing PCA. If center_data is false,
            it is expected that the data has already been centered and an exception
            will be thrown if it is not.

        Returns:
            Tuple[Tensor, Tensor]: The principal components (i.e., right singular vectors)
            and their corresponding singular values.

        Note: the algorithm assumes that the rows of X are the data points (after reshaping as needed).
        Consequently, the principal components, which are the eigenvectors of X.T @ X,
        are the right singular vectors in the SVD of X.
        """
        X_prime = self.prepare_data_forward(X, center_data=center_data)
        if self.low_rank:
            log(INFO, "Assuming data matrix is low rank, using low-rank PCA implementation.")
            m, n = X_prime.size(0), X_prime.size(1)
            if self.rank_estimation > m or self.rank_estimation > n:
                log(WARNING, "Estimate of data rank given by user is larger than the actual rank.")
            q = min(self.rank_estimation, m, n)
            _, singular_values, principal_components = torch.pca_lowrank(X_prime, q=q, center=False)
        else:
            if self.full_svd:
                log(INFO, "Performing full SVD on data matrix.")
            else:
                log(INFO, "Performing reduced SVD on data matrix.")
            U, s, Vh = torch.svd(X_prime, some=self.full_svd)
            principal_components = Vh.T
        return principal_components, singular_values
//  [python]
def forward(self, X: Tensor, center_data: bool) -> Tuple[Tensor, Tensor]:
        """
        Perform PCA on the data matrix X by computing its SVD.

        Args:
            X (Tensor): Data matrix.
            center_data (bool): If true, then the data mean will be subtracted
            from all data points prior to performing PCA. If center_data is false,
            it is expected that the data has already been centered and an exception
            will be thrown if it is not.

        Returns:
            Tuple[Tensor, Tensor]: The principal components (i.e., right singular vectors)
            and their corresponding singular values.

        Note: the algorithm assumes that the rows of X are the data points (after reshaping as needed).
        Consequently, the principal components, which are the eigenvectors of X.T @ X,
        are the right singular vectors in the SVD of X.
        """
        X_prime = self.prepare_data_forward(X, center_data=center_data)
        if self.low_rank:
            log(INFO, "Assuming data matrix is low rank, using low-rank PCA implementation.")
            m, n = X_prime.size(0), X_prime.size(1)
            if self.rank_estimation > m or self.rank_estimation > n:
                log(WARNING, "Estimate of data rank given by user is larger than the actual rank.")
            q = min(self.rank_estimation, m, n)
            _, singular_values, principal_components = torch.pca_lowrank(X_prime, q=q, center=False)
        else:
            if self.full_svd:
                log(INFO, "Performing full SVD on data matrix.")
            else:
                log(INFO, "Performing reduced SVD on data matrix.")
            _, singular_values, Vh = torch.linalg.svd(X_prime, full_matrices=self.full_svd)
            principal_components = Vh.T
        return principal_components, singular_values
//  [python]
def __init__(self, Y, r=0.99):
        # centerizing data
        self.mean = torch.mean(Y, axis=0)
        Y = Y - self.mean
        
        U, S, Vh = torch.svd(Y)
        cumuEnergy = S.cumsum(dim=0) / S.sum(dim=0)
        
        if r >= 1:
            rank = r 
        if r < 1:
            rank = (cumuEnergy>r).nonzero()[0][0]

        self.rank = rank
        # truncate the singular values and vector 
        U = U[:,0:rank]
        S = S[0:rank]
        Vh = Vh[0:rank,:]
        
        self.U = U
        self.S = S
        self.Vh = Vh
        
        self.Z = U @ S.diag_embed()
//  [python]
def __init__(self, Y, r=0.99):
        # centerizing data
        self.mean = torch.mean(Y, axis=0)
        Y = Y - self.mean
        
        U, S, Vh = torch.linalg.svd(Y, full_matrices=True)
        cumuEnergy = S.cumsum(dim=0) / S.sum(dim=0)
        
        if r >= 1:
            rank = r 
        if r < 1:
            rank = (cumuEnergy>r).nonzero()[0][0]

        self.rank = rank
        # truncate the singular values and vector 
        U = U[:,0:rank]
        S = S[0:rank]
        Vh = Vh[0:rank,:]
        
        self.U = U
        self.S = S
        self.Vh = Vh
        
        self.Z = U @ S.diag_embed()
//  [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() or has_hipsolver():
                backends.append("cusolver")

        ns = (12, 4, 2, 0)
        batches = ((), (0,), (1,), (2,), (2, 1), (0, 2))
        drivers = (None, 'gesvd', 'gesvdj', 'gesvda')

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

            for batch, m, n, driver in product(batches, ns, ns, drivers):
                if not (backend == 'cusolver' or driver is None):
                    # only test cases below and skip otherwise:
                    # - backend == 'cusolver' (driver can be anything)
                    # - backend != 'cusolver' (driver should only be None)
                    continue

                shape = batch + (m, n)
                k = min(m, n)
                A = make_arg(shape)
                U, S, Vh = torch.svd(A, full_matrices=False, driver=driver)
                self.assertEqual((U @ S.to(A.dtype).diag_embed()) @ Vh, A)

                U_f, S_f, Vh_f = torch.svd(A, full_matrices=True, driver=driver)
                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() or has_hipsolver():
                backends.append("cusolver")

        ns = (12, 4, 2, 0)
        batches = ((), (0,), (1,), (2,), (2, 1), (0, 2))
        drivers = (None, 'gesvd', 'gesvdj', 'gesvda')

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

            for batch, m, n, driver in product(batches, ns, ns, drivers):
                if not (backend == 'cusolver' or driver is None):
                    # only test cases below and skip otherwise:
                    # - backend == 'cusolver' (driver can be anything)
                    # - backend != 'cusolver' (driver should only be None)
                    continue

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

                U_f, S_f, Vh_f = torch.linalg.svd(A, full_matrices=True, driver=driver)
                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, driver=driver)
                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 analytical_mode(pred, fn_type):
    device = pred.device
    if CPUSVD:
        pred = pred.cpu()

    U, S, VT = torch.svd(pred)
    with torch.no_grad():  # sign can only change if the 3rd component of the svd is 0, then the sign does not matter
        s3sign = torch.det(torch.matmul(U, VT))
    diag = torch.stack((torch.ones_like(s3sign), torch.ones_like(s3sign), s3sign), -1)
    diag = torch.diag_embed(diag)

    pred_orth = U @ diag @ VT
    pred_orth = pred_orth.to(device)
    return pred_orth, s3sign
//  [python]
def analytical_mode(pred, fn_type):
    device = pred.device
    if CPUSVD:
        pred = pred.cpu()

    U, S, VT = torch.linalg.svd(pred)
    with torch.no_grad():  # sign can only change if the 3rd component of the svd is 0, then the sign does not matter
        s3sign = torch.det(torch.matmul(U, VT))
    diag = torch.stack((torch.ones_like(s3sign), torch.ones_like(s3sign), s3sign), -1)
    diag = torch.diag_embed(diag)

    pred_orth = U @ diag @ VT
    pred_orth = pred_orth.to(device)
    return pred_orth, s3sign
//  [python]
def project2orthogonal(matrix: torch.Tensor, rank:int, compute_device:torch.device):
	dim_1, dim_2 = matrix.shape[-2], matrix.shape[-1]
	if rank is None: rank = min(matrix.shape[-2:])
	try:
		if matrix.shape[-2] / matrix.shape[-1] >= 0.5:
			try:
				U, S, Vh = torch.svd(matrix, some=False)
				final =  U[..., :rank].to(compute_device) @ Vh[..., :rank, :].to(compute_device)
				if torch.any(torch.isnan(final)) or torch.any(torch.isinf(final)):
					U, S, Vh = torch.svd(matrix, some=False)
					final = U[..., :rank].to(compute_device) @ Vh[..., :rank, :].to(compute_device)
			except Exception as e:
				U, S, Vh = torch.svd(matrix, some=False)
				final = U[..., :rank].to(compute_device) @ Vh[..., :rank, :].to(compute_device)
		else:
			U, S, Vh = torch.svd(matrix, some=False)
			final = U[..., :rank].to(compute_device) @ Vh[..., :rank, :].to(compute_device)
		# a = U[..., :rank].to(compute_device) @ Vh[..., :rank, :].to(compute_device)
		# if torch.sum(torch.isnan(a)) > 0:
		# 	print("a", a, torch.sum(torch.isnan(a)), a.shape)
		if torch.any(torch.isnan(final)):
			print("gesvd & default failed")
			raise BaseException
		return final, S[..., :rank]
	except Exception as e:
		print(f'error {e}. using eigh, shape = {matrix.shape}')

		kk = 1e-2
		U, S, V, Vh = None, None, None, None
		mode = dim_2 > dim_1
		X = matrix @ matrix.transpose(-1, -2) if mode else matrix.transpose(-1, -2) @ matrix
		t = X.diagonal(dim1=-1, dim2=-2)
		t += kk
		del t
		eigvals, eigvecs = torch.linalg.eigh(X)
		eigvecs = eigvecs[..., -rank:].flip(-1)
		if mode:
			U = eigvecs
			Vh = U.transpose(-1, -2) @ matrix
			V, R = torch.linalg.qr(Vh.transpose(-1, -2))
			V = V.mul_(R.diagonal(dim1=-1, dim2=-2).sign()[..., None, :])
		else:
			V = eigvecs
			U = matrix @ V
			U, R = torch.linalg.qr(U)
			U = U.mul_(R.diagonal(dim1=-1, dim2=-2).sign()[..., None, :])
		UVh = U @ V.transpose(-1, -2)
		assert (U.transpose(-1, -2) @ U - torch.eye(U.shape[-1], device=U.device)).max().item() < 1e-4
		assert (V.transpose(-1, -2) @ V - torch.eye(V.shape[-1], device=V.device)).max().item() < 1e-4
		return UVh, S
//  [python]
def project2orthogonal(matrix: torch.Tensor, rank:int, compute_device:torch.device):
	dim_1, dim_2 = matrix.shape[-2], matrix.shape[-1]
	if rank is None: rank = min(matrix.shape[-2:])
	try:
		if matrix.shape[-2] / matrix.shape[-1] >= 0.5:
			try:
				U, S, Vh = torch.linalg.svd(matrix, full_matrices=False, driver='gesvda')
				final =  U[..., :rank].to(compute_device) @ Vh[..., :rank, :].to(compute_device)
				if torch.any(torch.isnan(final)) or torch.any(torch.isinf(final)):
					U, S, Vh = torch.linalg.svd(matrix, full_matrices=False)
					final = U[..., :rank].to(compute_device) @ Vh[..., :rank, :].to(compute_device)
			except Exception as e:
				U, S, Vh = torch.linalg.svd(matrix, full_matrices=False)
				final = U[..., :rank].to(compute_device) @ Vh[..., :rank, :].to(compute_device)
		else:
			U, S, Vh = torch.linalg.svd(matrix, full_matrices=False)
			final = U[..., :rank].to(compute_device) @ Vh[..., :rank, :].to(compute_device)
		# a = U[..., :rank].to(compute_device) @ Vh[..., :rank, :].to(compute_device)
		# if torch.sum(torch.isnan(a)) > 0:
		# 	print("a", a, torch.sum(torch.isnan(a)), a.shape)
		if torch.any(torch.isnan(final)):
			print("gesvd & default failed")
			raise BaseException
		return final, S[..., :rank]
	except Exception as e:
		print(f'error {e}. using eigh, shape = {matrix.shape}')

		kk = 1e-2
		U, S, V, Vh = None, None, None, None
		mode = dim_2 > dim_1
		X = matrix @ matrix.transpose(-1, -2) if mode else matrix.transpose(-1, -2) @ matrix
		t = X.diagonal(dim1=-1, dim2=-2)
		t += kk
		del t
		eigvals, eigvecs = torch.linalg.eigh(X)
		eigvecs = eigvecs[..., -rank:].flip(-1)
		if mode:
			U = eigvecs
			Vh = U.transpose(-1, -2) @ matrix
			V, R = torch.linalg.qr(Vh.transpose(-1, -2))
			V = V.mul_(R.diagonal(dim1=-1, dim2=-2).sign()[..., None, :])
		else:
			V = eigvecs
			U = matrix @ V
			U, R = torch.linalg.qr(U)
			U = U.mul_(R.diagonal(dim1=-1, dim2=-2).sign()[..., None, :])
		UVh = U @ V.transpose(-1, -2)
		assert (U.transpose(-1, -2) @ U - torch.eye(U.shape[-1], device=U.device)).max().item() < 1e-4
		assert (V.transpose(-1, -2) @ V - torch.eye(V.shape[-1], device=V.device)).max().item() < 1e-4
		return UVh, S
//  [python]
def low_rank_decomposition(weight, rank_ratio=0.1, parameter_ratio=0.15,
                           remove_criteria='max_eigenvalue',
                           log_level='INFO',
                           return_dict=False):
    """
    :param          weight: The matrix to decompose, of shape (H, W)
    :param      rank_ratio: rank_of_decomposed_matrix / rank_of_input_weight
    :param parameter_ratio: parameter_num_of_decomposed_matrix / (H * W). If specify, override rank_ratio
    :param remove_criteria: choose from ['max_eigenvalue', 'random', 'min_eigenvalue']
    :param       log_level: choose from ['IGNORE', 'INFO', 'DEBUG']
    :param     return_dict: Return a dict if True, else return a tuple (L, R)
    :return:
    """

    """parameter_ratio = rank * (H + W) / (H * W)"""
    """rank_ratio = """
    matrix_dimension = len(weight.size())
    assert matrix_dimension == 2, "Only Support 2D matrix"
    H, W = weight.size()

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

    if parameter_ratio is not None:
        reduced_rank = math.ceil(parameter_ratio * (H * W) / (H + W))
    else:
        reduced_rank = math.ceil(rank * rank_ratio)

    if remove_criteria == 'max_eigenvalue':
        L = U @ (torch.sqrt(torch.diag(S)[:, 0:reduced_rank]))
        R = torch.sqrt(torch.diag(S)[0:reduced_rank, :]) @ Vh
    elif remove_criteria == 'random':
        selected_index = random.choices(range(len(S)), k=reduced_rank)
        L = U @ (torch.sqrt(torch.diag(S)[:, selected_index]))
        R = torch.sqrt(torch.diag(S)[selected_index, :]) @ Vh
    elif remove_criteria == 'min_eigenvalue':
        len_s = len(S)
        L = U @ (torch.sqrt(torch.diag(S)[:, len_s - reduced_rank:]))
        R = torch.sqrt(torch.diag(S)[len_s - reduced_rank:, :]) @ Vh
    else:
        raise NameError("remove criteria not support")

    #########
    #  LOG  #
    #########
    if log_level == 'INFO':
        if not is_full_rank:
            print(f"It is not a full rank matrix. Rank: {rank} | H x W: {H}, {W}")
        print(f"Reduced Rank: {reduced_rank} | Num Parameters: {(H + W) * reduced_rank}")
    if log_level == 'DEBUG':
        print(f"W: ({H},{W}) | Rank: {rank} | U:{U.shape} | S:{S.shape} | Vh:{Vh.shape}")
        print(f"Reduced Rank: {reduced_rank} | Num Parameters: {(H + W) * reduced_rank}")
        print(f"L: {L.shape} | R: {R.shape}")

    if return_dict:
        return {"L": L, "R": R, "U": U, "S": S, "Vh": Vh, 'reduced_rank': reduced_rank}
    else:
        return L, R
//  [python]
def low_rank_decomposition(weight, rank_ratio=0.1, parameter_ratio=0.15,
                           remove_criteria='max_eigenvalue',
                           log_level='INFO',
                           return_dict=False):
    """
    :param          weight: The matrix to decompose, of shape (H, W)
    :param      rank_ratio: rank_of_decomposed_matrix / rank_of_input_weight
    :param parameter_ratio: parameter_num_of_decomposed_matrix / (H * W). If specify, override rank_ratio
    :param remove_criteria: choose from ['max_eigenvalue', 'random', 'min_eigenvalue']
    :param       log_level: choose from ['IGNORE', 'INFO', 'DEBUG']
    :param     return_dict: Return a dict if True, else return a tuple (L, R)
    :return:
    """

    """parameter_ratio = rank * (H + W) / (H * W)"""
    """rank_ratio = """
    matrix_dimension = len(weight.size())
    assert matrix_dimension == 2, "Only Support 2D matrix"
    H, W = weight.size()

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

    if parameter_ratio is not None:
        reduced_rank = math.ceil(parameter_ratio * (H * W) / (H + W))
    else:
        reduced_rank = math.ceil(rank * rank_ratio)

    if remove_criteria == 'max_eigenvalue':
        L = U @ (torch.sqrt(torch.diag(S)[:, 0:reduced_rank]))
        R = torch.sqrt(torch.diag(S)[0:reduced_rank, :]) @ Vh
    elif remove_criteria == 'random':
        selected_index = random.choices(range(len(S)), k=reduced_rank)
        L = U @ (torch.sqrt(torch.diag(S)[:, selected_index]))
        R = torch.sqrt(torch.diag(S)[selected_index, :]) @ Vh
    elif remove_criteria == 'min_eigenvalue':
        len_s = len(S)
        L = U @ (torch.sqrt(torch.diag(S)[:, len_s - reduced_rank:]))
        R = torch.sqrt(torch.diag(S)[len_s - reduced_rank:, :]) @ Vh
    else:
        raise NameError("remove criteria not support")

    #########
    #  LOG  #
    #########
    if log_level == 'INFO':
        if not is_full_rank:
            print(f"It is not a full rank matrix. Rank: {rank} | H x W: {H}, {W}")
        print(f"Reduced Rank: {reduced_rank} | Num Parameters: {(H + W) * reduced_rank}")
    if log_level == 'DEBUG':
        print(f"W: ({H},{W}) | Rank: {rank} | U:{U.shape} | S:{S.shape} | Vh:{Vh.shape}")
        print(f"Reduced Rank: {reduced_rank} | Num Parameters: {(H + W) * reduced_rank}")
        print(f"L: {L.shape} | R: {R.shape}")

    if return_dict:
        return {"L": L, "R": R, "U": U, "S": S, "Vh": Vh, 'reduced_rank': reduced_rank}
    else:
        return L, R
//  [python]
def fit(
        self,
        x: torch.Tensor
    ) -> Self:
        """Fit the Nystroem approximation.

        Parameters
        ----------
        x
            `torch.Tensor` of shape ``(n, d)`` where ``n`` is the number of samples and ``d`` is the dimensionality of
            the data.
        """
        n = len(x)
        n_components = n if self.n_components is None else self.n_components
        inds = torch.randperm(n)[:n_components]
        self.z = x[inds]
        K_zz = self.kernel(self.z,  self.z)
        K_zz += 1e-16 + torch.eye(n_components, device=K_zz.device)
        U, S, V = torch.svd(K_zz)
        self.K_zz_root_inv = (U / S.sqrt()) @ V
        return self
//  [python]
def fit(
        self,
        x: torch.Tensor
    ) -> Self:
        """Fit the Nystroem approximation.

        Parameters
        ----------
        x
            `torch.Tensor` of shape ``(n, d)`` where ``n`` is the number of samples and ``d`` is the dimensionality of
            the data.
        """
        n = len(x)
        n_components = n if self.n_components is None else self.n_components
        inds = torch.randperm(n)[:n_components]
        self.z = x[inds]
        K_zz = self.kernel(self.z,  self.z)
        K_zz += 1e-16 + torch.eye(n_components, device=K_zz.device)
        U, S, V = torch.linalg.svd(K_zz)
        self.K_zz_root_inv = (U / S.sqrt()) @ V
        return self
//  [python]
def main():
    parser = argparse.ArgumentParser(description='Model Stats Demo')
    parser.add_argument('--use-hf', default=False, action='store_true',
                        help='Use RADIO from HuggingFace Hub'
    )
    parser.add_argument('-v', '--model-version', default='radio_v2',
                        help='Which radio model to load.'
    )
    parser.add_argument('-d', '--dataset', default='food101',
                        help='The name of the dataset to classify'
    )

    args, _ = parser.parse_known_args()

    print('Loading model...')
    if args.use_hf:
        from transformers import AutoModel
        model: nn.Module = AutoModel.from_pretrained(f"nvidia/{args.model_version}", trust_remote_code=True)
    else:
        model: nn.Module = torch.hub.load('NVlabs/RADIO', 'radio_model', version=args.model_version, progress=True)

    model.cuda().eval()
    print('Done')

    spect_norms = []
    names = []

    for prm_name, prm in model.named_parameters():
        if prm.ndim < 2 or 'pos_embed' in prm_name:
            continue

        def _add_norm(n, p):
            w = p.flatten(1)
            U, S, V = torch.svd(w)

            spect_norms.append(S[0].clone())
            names.append(n)

        if prm_name.endswith('qkv.weight'):
            for n, p in zip(['q', 'k', 'v'], prm.chunk(3)):
                _add_norm(prm_name.replace('qkv', n), p)
        else:
            _add_norm(prm_name, prm)

    spect_norms = torch.stack(spect_norms)

    max_norm, max_norm_idx = spect_norms.max(dim=0)
    med_norm = spect_norms.median().item()
    avg_norm = spect_norms.mean().item()

    print(f'Spectral - Max: {max_norm.item()}, Median: {med_norm}, Mean: {avg_norm}')
    print(f'Max Parameter: {names[max_norm_idx.item()]}')
    sorted_norms, sorted_idxs = torch.sort(spect_norms, descending=True)
    print(f'Norms: {sorted_norms.cpu()}')

    for n, mod in model.named_modules():
        if isinstance(mod, Attention):
            mod.register_forward_hook(_attn_hook(n))
        elif isinstance(mod, Mlp):
            mod.register_forward_hook(_mlp_hook(n))

    transform = transforms.Compose([
        transforms.Resize([378, 378]),
        transforms.ToImage(),
        transforms.ToDtype(torch.float32, scale=True),
    ])

    ds_builder = load_dataset_builder(args.dataset, trust_remote_code=True)
    ds_builder.download_and_prepare()

    dataset = ds_builder.as_dataset(split='train').to_iterable_dataset(num_shards=8)
    dataset = dataset.map(lambda ex: dict(image=transform(ex['image'])))
    loader = DataLoader(dataset, batch_size=16, pin_memory=True, num_workers=8)

    n_iter = 100
    for i, batch in tqdm(enumerate(loader), total=n_iter):
        if i == n_iter:
            break

        images = batch['image'].cuda()

        model(images)

    _normalize_stats(n_iter)

    for layer_name, stats in LAYER_STATS.items():
        print(layer_name)
        for metric_name, metric in stats.items():
            print(f'\t{metric_name} - {metric}')
//  [python]
def main():
    parser = argparse.ArgumentParser(description='Model Stats Demo')
    parser.add_argument('--use-hf', default=False, action='store_true',
                        help='Use RADIO from HuggingFace Hub'
    )
    parser.add_argument('-v', '--model-version', default='radio_v2',
                        help='Which radio model to load.'
    )
    parser.add_argument('-d', '--dataset', default='food101',
                        help='The name of the dataset to classify'
    )

    args, _ = parser.parse_known_args()

    print('Loading model...')
    if args.use_hf:
        from transformers import AutoModel
        model: nn.Module = AutoModel.from_pretrained(f"nvidia/{args.model_version}", trust_remote_code=True)
    else:
        model: nn.Module = torch.hub.load('NVlabs/RADIO', 'radio_model', version=args.model_version, progress=True)

    model.cuda().eval()
    print('Done')

    spect_norms = []
    names = []

    for prm_name, prm in model.named_parameters():
        if prm.ndim < 2 or 'pos_embed' in prm_name:
            continue

        def _add_norm(n, p):
            w = p.flatten(1)
            _, s, _ = torch.linalg.svd(w)

            spect_norms.append(s[0].clone())
            names.append(n)

        if prm_name.endswith('qkv.weight'):
            for n, p in zip(['q', 'k', 'v'], prm.chunk(3)):
                _add_norm(prm_name.replace('qkv', n), p)
        else:
            _add_norm(prm_name, prm)

    spect_norms = torch.stack(spect_norms)

    max_norm, max_norm_idx = spect_norms.max(dim=0)
    med_norm = spect_norms.median().item()
    avg_norm = spect_norms.mean().item()

    print(f'Spectral - Max: {max_norm.item()}, Median: {med_norm}, Mean: {avg_norm}')
    print(f'Max Parameter: {names[max_norm_idx.item()]}')
    sorted_norms, sorted_idxs = torch.sort(spect_norms, descending=True)
    print(f'Norms: {sorted_norms.cpu()}')

    for n, mod in model.named_modules():
        if isinstance(mod, Attention):
            mod.register_forward_hook(_attn_hook(n))
        elif isinstance(mod, Mlp):
            mod.register_forward_hook(_mlp_hook(n))

    transform = transforms.Compose([
        transforms.Resize([378, 378]),
        transforms.ToImage(),
        transforms.ToDtype(torch.float32, scale=True),
    ])

    ds_builder = load_dataset_builder(args.dataset, trust_remote_code=True)
    ds_builder.download_and_prepare()

    dataset = ds_builder.as_dataset(split='train').to_iterable_dataset(num_shards=8)
    dataset = dataset.map(lambda ex: dict(image=transform(ex['image'])))
    loader = DataLoader(dataset, batch_size=16, pin_memory=True, num_workers=8)

    n_iter = 100
    for i, batch in tqdm(enumerate(loader), total=n_iter):
        if i == n_iter:
            break

        images = batch['image'].cuda()

        model(images)

    _normalize_stats(n_iter)

    for layer_name, stats in LAYER_STATS.items():
        print(layer_name)
        for metric_name, metric in stats.items():
            print(f'\t{metric_name} - {metric}')
//  [python]
def align_pcl(Y, X, weight=None, fixed_scale=False):
    """align similarity transform to align X with Y using umeyama method
    X' = s * R * X + t is aligned with Y
    :param Y (*, N, 3) first trajectory
    :param X (*, N, 3) second trajectory
    :param weight (*, N, 1) optional weight of valid correspondences
    :returns s (*, 1), R (*, 3, 3), t (*, 3)
    """
    *dims, N, _ = Y.shape
    N = torch.ones(*dims, 1, 1) * N

    if weight is not None:
        Y = Y * weight
        X = X * weight
        N = weight.sum(dim=-2, keepdim=True)  # (*, 1, 1)

    # subtract mean
    my = Y.sum(dim=-2) / N[..., 0]  # (*, 3)
    mx = X.sum(dim=-2) / N[..., 0]
    y0 = Y - my[..., None, :]  # (*, N, 3)
    x0 = X - mx[..., None, :]

    if weight is not None:
        y0 = y0 * weight
        x0 = x0 * weight

    # correlation
    C = torch.matmul(y0.transpose(-1, -2), x0) / N  # (*, 3, 3)
    U, D, Vh = torch.svd(C)  # (*, 3, 3), (*, 3), (*, 3, 3)

    S = torch.eye(3).reshape(*(1,) * (len(dims)), 3, 3).repeat(*dims, 1, 1)
    neg = torch.det(U) * torch.det(Vh.transpose(-1, -2)) < 0
    S[neg, 2, 2] = -1

    R = torch.matmul(U, torch.matmul(S, Vh))  # (*, 3, 3)

    D = torch.diag_embed(D)  # (*, 3, 3)
    if fixed_scale:
        s = torch.ones(*dims, 1, device=Y.device, dtype=torch.float32)
    else:
        var = torch.sum(torch.square(x0), dim=(-1, -2), keepdim=True) / N  # (*, 1, 1)
        s = (
            torch.diagonal(torch.matmul(D, S), dim1=-2, dim2=-1).sum(
                dim=-1, keepdim=True
            )
            / var[..., 0]
        )  # (*, 1)

    t = my - s * torch.matmul(R, mx[..., None])[..., 0]  # (*, 3)

    return s, R, t
//  [python]
def align_pcl(Y, X, weight=None, fixed_scale=False):
    """align similarity transform to align X with Y using umeyama method
    X' = s * R * X + t is aligned with Y
    :param Y (*, N, 3) first trajectory
    :param X (*, N, 3) second trajectory
    :param weight (*, N, 1) optional weight of valid correspondences
    :returns s (*, 1), R (*, 3, 3), t (*, 3)
    """
    *dims, N, _ = Y.shape
    N = torch.ones(*dims, 1, 1) * N

    if weight is not None:
        Y = Y * weight
        X = X * weight
        N = weight.sum(dim=-2, keepdim=True)  # (*, 1, 1)

    # subtract mean
    my = Y.sum(dim=-2) / N[..., 0]  # (*, 3)
    mx = X.sum(dim=-2) / N[..., 0]
    y0 = Y - my[..., None, :]  # (*, N, 3)
    x0 = X - mx[..., None, :]

    if weight is not None:
        y0 = y0 * weight
        x0 = x0 * weight

    # correlation
    C = torch.matmul(y0.transpose(-1, -2), x0) / N  # (*, 3, 3)
    U, D, Vh = torch.linalg.svd(C)  # (*, 3, 3), (*, 3), (*, 3, 3)

    S = torch.eye(3).reshape(*(1,) * (len(dims)), 3, 3).repeat(*dims, 1, 1)
    neg = torch.det(U) * torch.det(Vh.transpose(-1, -2)) < 0
    S[neg, 2, 2] = -1

    R = torch.matmul(U, torch.matmul(S, Vh))  # (*, 3, 3)

    D = torch.diag_embed(D)  # (*, 3, 3)
    if fixed_scale:
        s = torch.ones(*dims, 1, device=Y.device, dtype=torch.float32)
    else:
        var = torch.sum(torch.square(x0), dim=(-1, -2), keepdim=True) / N  # (*, 1, 1)
        s = (
            torch.diagonal(torch.matmul(D, S), dim1=-2, dim2=-1).sum(
                dim=-1, keepdim=True
            )
            / var[..., 0]
        )  # (*, 1)

    t = my - s * torch.matmul(R, mx[..., None])[..., 0]  # (*, 3)

    return s, R, t
//  [python]
def postprocess(self, net: nn.Module, data: Any):
        inputs = data.cuda()

        # Logit of Block 4 feature
        feat1 = net.intermediate_forward(inputs, layer_index=4)
        B, C, H, W = feat1.size()
        feat1 = feat1.view(B, C, H * W)
        if self.args.accelerate:
            feat1 = feat1 - power_iteration(feat1, iter=20)
        else:
            u, s, v = torch.svd(feat1)
            feat1 = feat1 - s[:, 0:1].unsqueeze(2) * u[:, :, 0:1].bmm(
                v[:, 0:1, :])
        feat1 = feat1.view(B, C, H, W)
        logits1 = net.fc(torch.flatten(net.avgpool(feat1), 1))

        # Logit of Block 3 feature
        feat2 = net.intermediate_forward(inputs, layer_index=3)
        B, C, H, W = feat2.size()
        feat2 = feat2.view(B, C, H * W)
        if self.args.accelerate:
            feat2 = feat2 - power_iteration(feat2, iter=20)
        else:
            u, s, v = torch.svd(feat2)
            feat2 = feat2 - s[:, 0:1].unsqueeze(2) * u[:, :, 0:1].bmm(
                v[:, 0:1, :])
        feat2 = feat2.view(B, C, H, W)
        feat2 = net.layer4(feat2)
        logits2 = net.fc(torch.flatten(net.avgpool(feat2), 1))

        # Fusion at the logit space
        logits = (logits1 + logits2) / 2
        conf = self.args.temperature * torch.logsumexp(
            logits / self.args.temperature, dim=1)

        _, pred = torch.max(logits, dim=1)
        return pred, conf
//  [python]
def postprocess(self, net: nn.Module, data: Any):
        inputs = data.cuda()

        # Logit of Block 4 feature
        feat1 = net.intermediate_forward(inputs, layer_index=4)
        B, C, H, W = feat1.size()
        feat1 = feat1.view(B, C, H * W)
        if self.args.accelerate:
            feat1 = feat1 - power_iteration(feat1, iter=20)
        else:
            u, s, v = torch.linalg.svd(feat1, full_matrices=False)
            feat1 = feat1 - s[:, 0:1].unsqueeze(2) * u[:, :, 0:1].bmm(
                v[:, 0:1, :])
        feat1 = feat1.view(B, C, H, W)
        logits1 = net.fc(torch.flatten(net.avgpool(feat1), 1))

        # Logit of Block 3 feature
        feat2 = net.intermediate_forward(inputs, layer_index=3)
        B, C, H, W = feat2.size()
        feat2 = feat2.view(B, C, H * W)
        if self.args.accelerate:
            feat2 = feat2 - power_iteration(feat2, iter=20)
        else:
            u, s, v = torch.linalg.svd(feat2, full_matrices=False)
            feat2 = feat2 - s[:, 0:1].unsqueeze(2) * u[:, :, 0:1].bmm(
                v[:, 0:1, :])
        feat2 = feat2.view(B, C, H, W)
        feat2 = net.layer4(feat2)
        logits2 = net.fc(torch.flatten(net.avgpool(feat2), 1))

        # Fusion at the logit space
        logits = (logits1 + logits2) / 2
        conf = self.args.temperature * torch.logsumexp(
            logits / self.args.temperature, dim=1)

        _, pred = torch.max(logits, dim=1)
        return pred, conf
//  [python]
def batch_find_rigid_alignment(A, B):
    """
        Input:
            A:  (B, 3)
            B:  (B, 3)
        Returns:
            Rs: (B, 3, 3)
            Ts: (B, 3)
    """
    assert len(A.shape) == 3
    a_mean = A.mean(axis=1, keepdim=True)
    b_mean = B.mean(axis=1, keepdim=True)
    A_c = A - a_mean
    B_c = B - b_mean
    # Covariance matrix
    H = einops.einsum(A_c, B_c, "b n c1, b n c2 -> b c1 c2")

    # The returned V in torch.linalg.svd is different from torch.svd
    #   tensor.mT means tensor.transpose(-1, -2)
    #   tensor.T means tensor.transpose(-1, -2, -3, ...)
    U, S, V = torch.svd(H)
    VmT = V.mT
    V = VmT.mT
    # Rotation matrix
    R = einops.einsum(V, U.transpose(2, 1), "b c1 c2, b c2 c3 -> b c1 c3")
    # Translation vector
    t = b_mean - einops.einsum(R, a_mean, "b c1 c2, b n c2 -> b n c1")
    return R, t.squeeze()
//  [python]
def batch_find_rigid_alignment(A, B):
    """
        Input:
            A:  (B, 3)
            B:  (B, 3)
        Returns:
            Rs: (B, 3, 3)
            Ts: (B, 3)
    """
    assert len(A.shape) == 3
    a_mean = A.mean(axis=1, keepdim=True)
    b_mean = B.mean(axis=1, keepdim=True)
    A_c = A - a_mean
    B_c = B - b_mean
    # Covariance matrix
    H = einops.einsum(A_c, B_c, "b n c1, b n c2 -> b c1 c2")

    # The returned V in torch.linalg.svd is different from torch.svd
    #   tensor.mT means tensor.transpose(-1, -2)
    #   tensor.T means tensor.transpose(-1, -2, -3, ...)
    U, S, VmT = torch.linalg.svd(H)
    V = VmT.mT
    # Rotation matrix
    R = einops.einsum(V, U.transpose(2, 1), "b c1 c2, b c2 c3 -> b c1 c3")
    # Translation vector
    t = b_mean - einops.einsum(R, a_mean, "b c1 c2, b n c2 -> b n c1")
    return R, t.squeeze()
//  [python]
def get_angles(covar_mat, node_dir):
    if covar_mat is not None and node_dir is not None and node_dir.shape[0] != 0:
        non_zero_dir_idx = (node_dir == torch.zeros_like(node_dir[0])).sum(1) == 0
        covar_mat = covar_mat[non_zero_dir_idx]
        node_dir = node_dir[non_zero_dir_idx]
        if covar_mat.shape[0] == 0:
            return None, None, None, None

        if len(covar_mat.shape) == 3:
            u, s, vh = torch.svd(covar_mat)
            pred_dir = u[:, :, 0]
        else:
            # grad on pos
            pred_dir = covar_mat
            s = covar_mat

        true_dir = node_dir[:, :pred_dir.shape[1]]
        angles = torch.rad2deg(torch.arccos(F.cosine_similarity(pred_dir, true_dir, dim=1).clamp(-1+1e-8, +1-1e-8)))
        angles[angles > 90] = 180 - angles[angles > 90]
        avg_angles = angles.mean().reshape(1)

        avg_eigen_ratio = (s[:, 1] / s[:, 0]).mean().reshape(1)
    else:
        return None, None, None, None
    return avg_angles, angles, avg_eigen_ratio, (pred_dir, true_dir)
//  [python]
def get_angles(covar_mat, node_dir):
    if covar_mat is not None and node_dir is not None and node_dir.shape[0] != 0:
        non_zero_dir_idx = (node_dir == torch.zeros_like(node_dir[0])).sum(1) == 0
        covar_mat = covar_mat[non_zero_dir_idx]
        node_dir = node_dir[non_zero_dir_idx]
        if covar_mat.shape[0] == 0:
            return None, None, None, None

        if len(covar_mat.shape) == 3:
            u, s, vh = torch.linalg.svd(covar_mat)
            pred_dir = u[:, :, 0]
        else:
            # grad on pos
            pred_dir = covar_mat
            s = covar_mat

        true_dir = node_dir[:, :pred_dir.shape[1]]
        angles = torch.rad2deg(torch.arccos(F.cosine_similarity(pred_dir, true_dir, dim=1).clamp(-1+1e-8, +1-1e-8)))
        angles[angles > 90] = 180 - angles[angles > 90]
        avg_angles = angles.mean().reshape(1)

        avg_eigen_ratio = (s[:, 1] / s[:, 0]).mean().reshape(1)
    else:
        return None, None, None, None
    return avg_angles, angles, avg_eigen_ratio, (pred_dir, true_dir)
//  [python]
def find_alignment_transform(traces):
    centers = traces.mean(-2, keepdim=True)
    traces = traces - centers

    p1, p2 = traces[0], traces[1:]

    C = torch.einsum("i j k, j l -> i k l", p2, p1)
    U, S, V = torch.svd(C)
    U = torch.matmul(V, W)
    U = torch.matmul(torch.stack([torch.ones(len(p2), device=U.device),torch.ones(len(p2), device=U.device),torch.det(U)], dim=1)[:,:,None] * V,  W)

    return torch.cat([torch.eye(3, device=U.device)[None], U]), centers
//  [python]
def find_alignment_transform(traces):
    centers = traces.mean(-2, keepdim=True)
    traces = traces - centers

    p1, p2 = traces[0], traces[1:]

    C = torch.einsum("i j k, j l -> i k l", p2, p1)
    V, _, W = torch.linalg.svd(C)
    U = torch.matmul(V, W)
    U = torch.matmul(torch.stack([torch.ones(len(p2), device=U.device),torch.ones(len(p2), device=U.device),torch.linalg.det(U)], dim=1)[:,:,None] * V,  W)

    return torch.cat([torch.eye(3, device=U.device)[None], U]), centers
//  [python]
def rotation_from_svd(points1, points2, center1=None, center2=None):
    """Compute rotation matrix from two point clouds using SVD.

    Args:
        points1: A tensor of shape (..., npts, 3)
        points2: A tensor of shape (..., npts, 3)
        cetner1: A tensor of shape (..., 3) representing the center of points1.
        cetner2: A tensor of shape (..., 3) representing the center of points2.

    Returns:
        A tensor of shape (..., 3, 3) representing the rotation matrix.
    """
    if center1 is None:
        points1 = points1 - points1.mean(dim=-2, keepdim=True)
    else:
        points1 = points1 - center1.unsqueeze(-2)

    if center2 is None:
        points2 = points2 - points2.mean(dim=-2, keepdim=True)
    else:
        points2 = points2 - center2.unsqueeze(-2)

    # compute svd
    H = points2.transpose(-2, -1) @ points1
    U, S, Vh = torch.svd(H)
    V = Vh.transpose(-2, -1)
    R = V @ U.transpose(-2, -1)

    # if the determinant(R) < 0, multiply the 3rd column of V with -1
    inverse_V = torch.stack([
        V[..., 0], V[..., 1], -V[..., 2]
    ], dim=-1)
    V = torch.where(torch.det(R).unsqueeze(-1).unsqueeze(-1) < 0,
                    inverse_V, V)
    R = V @ U.transpose(-2, -1)

    return R
//  [python]
def rotation_from_svd(points1, points2, center1=None, center2=None):
    """Compute rotation matrix from two point clouds using SVD.

    Args:
        points1: A tensor of shape (..., npts, 3)
        points2: A tensor of shape (..., npts, 3)
        cetner1: A tensor of shape (..., 3) representing the center of points1.
        cetner2: A tensor of shape (..., 3) representing the center of points2.

    Returns:
        A tensor of shape (..., 3, 3) representing the rotation matrix.
    """
    if center1 is None:
        points1 = points1 - points1.mean(dim=-2, keepdim=True)
    else:
        points1 = points1 - center1.unsqueeze(-2)

    if center2 is None:
        points2 = points2 - points2.mean(dim=-2, keepdim=True)
    else:
        points2 = points2 - center2.unsqueeze(-2)

    # compute svd
    H = points2.transpose(-2, -1) @ points1
    U, S, Vh = torch.linalg.svd(H)
    V = Vh.transpose(-2, -1)
    R = V @ U.transpose(-2, -1)

    # if the determinant(R) < 0, multiply the 3rd column of V with -1
    inverse_V = torch.stack([
        V[..., 0], V[..., 1], -V[..., 2]
    ], dim=-1)
    V = torch.where(torch.linalg.det(R).unsqueeze(-1).unsqueeze(-1) < 0,
                    inverse_V, V)
    R = V @ U.transpose(-2, -1)

    return R
//  [python]
def get_orthogonal_matrix(weights, rank, type):
    module_params = weights

    if module_params.data.dtype != torch.float:
        float_data = False
        original_type = module_params.data.dtype
        original_device = module_params.data.device
        matrix = module_params.data.float()
    else:
        float_data = True
        matrix = module_params.data

    U, s, Vh = torch.svd(matrix)

    # make the smaller matrix always to be orthogonal matrix
    if type == "right":
        # A = U[:, :rank] @ torch.diag(s[:rank])
        B = Vh[:rank, :]

        if not float_data:
            B = B.to(original_device).type(original_type)
        return B
    elif type == "left":
        A = U[:, :rank]
        # B = torch.diag(s[:rank]) @ Vh[:rank, :]
        if not float_data:
            A = A.to(original_device).type(original_type)
        return A
    elif type == "full":
        A = U[:, :rank]
        B = Vh[:rank, :]
        if not float_data:
            A = A.to(original_device).type(original_type)
            B = B.to(original_device).type(original_type)
        return [A, B]
    else:
        raise ValueError("type should be left, right or full")
//  [python]
def get_orthogonal_matrix(weights, rank, type):
    module_params = weights

    if module_params.data.dtype != torch.float:
        float_data = False
        original_type = module_params.data.dtype
        original_device = module_params.data.device
        matrix = module_params.data.float()
    else:
        float_data = True
        matrix = module_params.data

    U, s, Vh = torch.linalg.svd(matrix, full_matrices=False)

    # make the smaller matrix always to be orthogonal matrix
    if type == "right":
        # A = U[:, :rank] @ torch.diag(s[:rank])
        B = Vh[:rank, :]

        if not float_data:
            B = B.to(original_device).type(original_type)
        return B
    elif type == "left":
        A = U[:, :rank]
        # B = torch.diag(s[:rank]) @ Vh[:rank, :]
        if not float_data:
            A = A.to(original_device).type(original_type)
        return A
    elif type == "full":
        A = U[:, :rank]
        B = Vh[:rank, :]
        if not float_data:
            A = A.to(original_device).type(original_type)
            B = B.to(original_device).type(original_type)
        return [A, B]
    else:
        raise ValueError("type should be left, right or full")