jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def tests():
        """
        Basic unit tests and test to check gradients
        """
        ndim = 2
        nb = 1
        # generate random base SPD matrix
        A = torch.randn((1,ndim,ndim), dtype=torch.double)
        U, s, _ = torch.svd(A)

        threshold = torch.tensor([1e-3], dtype=torch.double)

        # generate batches
        # linear case (all eigenvalues are above the threshold)
        s = threshold * 1e1 + torch.rand((nb,ndim), dtype=torch.double) * threshold
        M = U @ torch.diag_embed(s) @ U.transpose(-1,-2)

        assert (sym_reeig.apply(M, threshold, False).allclose(M))
        M.requires_grad_(True)
        assert(gradcheck(sym_reeig.apply, (M, threshold, True)))

        # non-linear case (some eigenvalues are below the threshold)
        s = torch.rand((nb,ndim), dtype=torch.double) * threshold
        s[::2] += threshold
        M = U @ torch.diag_embed(s) @ U.transpose(-1,-2)
        assert (~sym_reeig.apply(M, threshold, False).allclose(M))
        M.requires_grad_(True)
        assert(gradcheck(sym_reeig.apply, (M, threshold, True)))

        # linear case, all eigenvalues are identical
        s = torch.ones((nb,ndim), dtype=torch.double)
        M = U @ torch.diag_embed(s) @ U.transpose(-1,-2)
        assert (sym_reeig.apply(M, threshold, True).allclose(M))
        M.requires_grad_(True)
        assert(gradcheck(sym_reeig.apply, (M, threshold, True)))
//  [python]
def tests():
        """
        Basic unit tests and test to check gradients
        """
        ndim = 2
        nb = 1
        # generate random base SPD matrix
        A = torch.randn((1,ndim,ndim), dtype=torch.double)
        U, s, _ = torch.linalg.svd(A)

        threshold = torch.tensor([1e-3], dtype=torch.double)

        # generate batches
        # linear case (all eigenvalues are above the threshold)
        s = threshold * 1e1 + torch.rand((nb,ndim), dtype=torch.double) * threshold
        M = U @ torch.diag_embed(s) @ U.transpose(-1,-2)

        assert (sym_reeig.apply(M, threshold, False).allclose(M))
        M.requires_grad_(True)
        assert(gradcheck(sym_reeig.apply, (M, threshold, True)))

        # non-linear case (some eigenvalues are below the threshold)
        s = torch.rand((nb,ndim), dtype=torch.double) * threshold
        s[::2] += threshold
        M = U @ torch.diag_embed(s) @ U.transpose(-1,-2)
        assert (~sym_reeig.apply(M, threshold, False).allclose(M))
        M.requires_grad_(True)
        assert(gradcheck(sym_reeig.apply, (M, threshold, True)))

        # linear case, all eigenvalues are identical
        s = torch.ones((nb,ndim), dtype=torch.double)
        M = U @ torch.diag_embed(s) @ U.transpose(-1,-2)
        assert (sym_reeig.apply(M, threshold, True).allclose(M))
        M.requires_grad_(True)
        assert(gradcheck(sym_reeig.apply, (M, threshold, True)))
//  [python]
def random_fullrank_matrix_distinct_singular_value(matrix_size, *batch_dims,
                                                   **kwargs):
    dtype = kwargs.get('dtype', torch.double)
    device = kwargs.get('device', 'cpu')
    silent = kwargs.get("silent", False)
    if silent and not torch._C.has_lapack:
        return torch.ones(matrix_size, matrix_size, dtype=dtype, device=device)

    A = torch.randn(batch_dims + (matrix_size, matrix_size), dtype=dtype, device=device)
    u, _, vh = torch.svd(A, some=False)
    real_dtype = A.real.dtype if A.dtype.is_complex else A.dtype
    s = torch.arange(1., matrix_size + 1, dtype=real_dtype, device=device).mul_(1.0 / (matrix_size + 1))
    return (u * s.to(A.dtype)) @ vh
//  [python]
def random_fullrank_matrix_distinct_singular_value(matrix_size, *batch_dims,
                                                   **kwargs):
    dtype = kwargs.get('dtype', torch.double)
    device = kwargs.get('device', 'cpu')
    silent = kwargs.get("silent", False)
    if silent and not torch._C.has_lapack:
        return torch.ones(matrix_size, matrix_size, dtype=dtype, device=device)

    A = torch.randn(batch_dims + (matrix_size, matrix_size), dtype=dtype, device=device)
    u, _, vh = torch.linalg.svd(A, full_matrices=False)
    real_dtype = A.real.dtype if A.dtype.is_complex else A.dtype
    s = torch.arange(1., matrix_size + 1, dtype=real_dtype, device=device).mul_(1.0 / (matrix_size + 1))
    return (u * s.to(A.dtype)) @ vh
//  [python]
def random_matrix(rows, columns, *batch_dims, **kwargs):
    """Return rectangular matrix or batches of rectangular matrices.

    Parameters:
      dtype - the data type
      device - the device kind
      singular - when True, the output will be singular
    """
    dtype = kwargs.get('dtype', torch.double)
    device = kwargs.get('device', 'cpu')
    silent = kwargs.get("silent", False)
    singular = kwargs.get("singular", False)
    if silent and not torch._C.has_lapack:
        return torch.ones(rows, columns, dtype=dtype, device=device)

    A = torch.randn(batch_dims + (rows, columns), dtype=dtype, device=device)
    u, s, vh = torch.svd(A)
    k = min(rows, columns)
    s = torch.linspace(1 / (k + 1), 1, k, dtype=dtype, device=device)
    if singular:
        # make matrix singular
        s[k - 1] = 0
        if k > 2:
            # increase the order of singularity so that the pivoting
            # in LU factorization will be non-trivial
            s[0] = 0
    return (u * s.unsqueeze(-2)) @ vh
//  [python]
def random_matrix(rows, columns, *batch_dims, **kwargs):
    """Return rectangular matrix or batches of rectangular matrices.

    Parameters:
      dtype - the data type
      device - the device kind
      singular - when True, the output will be singular
    """
    dtype = kwargs.get('dtype', torch.double)
    device = kwargs.get('device', 'cpu')
    silent = kwargs.get("silent", False)
    singular = kwargs.get("singular", False)
    if silent and not torch._C.has_lapack:
        return torch.ones(rows, columns, dtype=dtype, device=device)

    A = torch.randn(batch_dims + (rows, columns), dtype=dtype, device=device)
    u, _, vh = torch.linalg.svd(A, full_matrices=False)
    k = min(rows, columns)
    s = torch.linspace(1 / (k + 1), 1, k, dtype=dtype, device=device)
    if singular:
        # make matrix singular
        s[k - 1] = 0
        if k > 2:
            # increase the order of singularity so that the pivoting
            # in LU factorization will be non-trivial
            s[0] = 0
    return (u * s.unsqueeze(-2)) @ vh
//  [python]
def extract_lora(diff, rank):
    conv2d = (len(diff.shape) == 4)
    kernel_size = None if not conv2d else diff.size()[2:4]
    conv2d_3x3 = conv2d and kernel_size != (1, 1)
    out_dim, in_dim = diff.size()[0:2]
    rank = min(rank, in_dim, out_dim)

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

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

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

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

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

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

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

    U = U.clamp(low_val, hi_val)
    Vh = Vh.clamp(low_val, hi_val)
    if conv2d:
        U = U.reshape(out_dim, rank, 1, 1)
        Vh = Vh.reshape(rank, in_dim, kernel_size[0], kernel_size[1])
    return (U, Vh)
//  [python]
def convert_to_standard_lora(self):
        # convert down and up matrix to linear layer, nn.Linear call F.linear
        with torch.no_grad():
            # merge two lora weight
            new_weight = torch.mm(self.up.weight, self.down.weight)
            # torch.svd
            U, s, Vt = torch.svd(new_weight)

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

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

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

            #  Remove abandoned layers and variables
            del self.down, self.up

            # Redefine down and up layers
            self.down = nn.Linear(self.in_features, r, bias=False)
            self.down.weight = down
            self.up = nn.Linear(r, self.out_features, bias=False)
            self.up.weight = nn.Parameter(up)
//  [python]
def convert_to_standard_lora(self):
        # convert down and up matrix to linear layer, nn.Linear call F.linear
        with torch.no_grad():
            # merge two lora weight
            new_weight = torch.mm(self.up.weight, self.down.weight)
            # torch.svd
            U, s, Vt = torch.linalg.svd(new_weight, full_matrices=False)

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

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

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

            #  Remove abandoned layers and variables
            del self.down, self.up

            # Redefine down and up layers
            self.down = nn.Linear(self.in_features, r, bias=False)
            self.down.weight = down
            self.up = nn.Linear(r, self.out_features, bias=False)
            self.up.weight = nn.Parameter(up)
//  [python]
def orthogonal_procrustes(self, A, B):
        u, s, vt = torch.svd(A.permute(0,2,1).bmm(B))
        return u.matmul(vt), s.sum(dim=1)
//  [python]
def orthogonal_procrustes(self, A, B):
        u, s, vt = torch.linalg.svd(A.permute(0,2,1).bmm(B))
        return u.matmul(vt), s.sum(dim=1)
//  [python]
def fit_transform(self, dataset: Tensor) -> Tensor:
        """Fit and transform PCA to dataset.

        Args:
          dataset (Tensor): Dataset to which the PCA if fit and transformed

        Returns:
          Transformed dataset
        """
        mean = dataset.mean(dim=0)
        dataset -= mean
        num_components = int(self.n_components)
        self.num_components = Tensor([num_components])

        u, s, v_h = torch.linalg.svd(dataset)
        self.singular_vectors = v_h.transpose(-2, -1)[:, :num_components]
        self.mean = mean

        return torch.matmul(dataset, self.singular_vectors)
//  [python]
def fit_transform(self, dataset: Tensor) -> Tensor:
        """Fit and transform PCA to dataset.

        Args:
          dataset (Tensor): Dataset to which the PCA if fit and transformed

        Returns:
          Transformed dataset
        """
        mean = dataset.mean(dim=0)
        dataset -= mean
        num_components = int(self.n_components)
        self.num_components = Tensor([num_components])

        v_h = torch.linalg.svd(dataset)[-1]
        self.singular_vectors = v_h.transpose(-2, -1)[:, :num_components]
        self.mean = mean

        return torch.matmul(dataset, self.singular_vectors)
//  [python]
def _compute_delta_joint_pos(self, delta_pose: torch.Tensor, jacobian: torch.Tensor) -> torch.Tensor:
        """Computes the change in joint position that yields the desired change in pose.

        The method uses the Jacobian mapping from joint-space velocities to end-effector velocities
        to compute the delta-change in the joint-space that moves the robot closer to a desired
        end-effector position.

        Args:
            delta_pose: The desired delta pose in shape (N, 3) or (N, 6).
            jacobian: The geometric jacobian matrix in shape (N, 3, num_joints) or (N, 6, num_joints).

        Returns:
            The desired delta in joint space. Shape is (N, num-jointsß).
        """
        if self.cfg.ik_params is None:
            raise RuntimeError(f"Inverse-kinematics parameters for method '{self.cfg.ik_method}' is not defined!")
        # compute the delta in joint-space
        if self.cfg.ik_method == "pinv":  # Jacobian pseudo-inverse
            # parameters
            k_val = self.cfg.ik_params["k_val"]
            # computation
            jacobian_pinv = torch.linalg.pinv(jacobian)
            delta_joint_pos = k_val * jacobian_pinv @ delta_pose.unsqueeze(-1)
            delta_joint_pos = delta_joint_pos.squeeze(-1)
        elif self.cfg.ik_method == "svd":  # adaptive SVD
            # parameters
            k_val = self.cfg.ik_params["k_val"]
            min_singular_value = self.cfg.ik_params["min_singular_value"]
            # computation
            # U: 6xd, S: dxd, V: d x num-joint
            U, S, Vh = torch.svd(jacobian)
            S_inv = 1.0 / S
            S_inv = torch.where(S > min_singular_value, S_inv, torch.zeros_like(S_inv))
            jacobian_pinv = (
                torch.transpose(Vh, dim0=1, dim1=2)[:, :, :6]
                @ torch.diag_embed(S_inv)
                @ torch.transpose(U, dim0=1, dim1=2)
            )
            delta_joint_pos = k_val * jacobian_pinv @ delta_pose.unsqueeze(-1)
            delta_joint_pos = delta_joint_pos.squeeze(-1)
        elif self.cfg.ik_method == "trans":  # Jacobian transpose
            # parameters
            k_val = self.cfg.ik_params["k_val"]
            # computation
            jacobian_T = torch.transpose(jacobian, dim0=1, dim1=2)
            delta_joint_pos = k_val * jacobian_T @ delta_pose.unsqueeze(-1)
            delta_joint_pos = delta_joint_pos.squeeze(-1)
        elif self.cfg.ik_method == "dls":  # damped least squares
            # parameters
            lambda_val = self.cfg.ik_params["lambda_val"]
            # computation
            jacobian_T = torch.transpose(jacobian, dim0=1, dim1=2)
            lambda_matrix = (lambda_val**2) * torch.eye(n=jacobian.shape[1], device=self._device)
            delta_joint_pos = (
                jacobian_T @ torch.inverse(jacobian @ jacobian_T + lambda_matrix) @ delta_pose.unsqueeze(-1)
            )
            delta_joint_pos = delta_joint_pos.squeeze(-1)
        else:
            raise ValueError(f"Unsupported inverse-kinematics method: {self.cfg.ik_method}")

        return delta_joint_pos
//  [python]
def _compute_delta_joint_pos(self, delta_pose: torch.Tensor, jacobian: torch.Tensor) -> torch.Tensor:
        """Computes the change in joint position that yields the desired change in pose.

        The method uses the Jacobian mapping from joint-space velocities to end-effector velocities
        to compute the delta-change in the joint-space that moves the robot closer to a desired
        end-effector position.

        Args:
            delta_pose: The desired delta pose in shape (N, 3) or (N, 6).
            jacobian: The geometric jacobian matrix in shape (N, 3, num_joints) or (N, 6, num_joints).

        Returns:
            The desired delta in joint space. Shape is (N, num-jointsß).
        """
        if self.cfg.ik_params is None:
            raise RuntimeError(f"Inverse-kinematics parameters for method '{self.cfg.ik_method}' is not defined!")
        # compute the delta in joint-space
        if self.cfg.ik_method == "pinv":  # Jacobian pseudo-inverse
            # parameters
            k_val = self.cfg.ik_params["k_val"]
            # computation
            jacobian_pinv = torch.linalg.pinv(jacobian)
            delta_joint_pos = k_val * jacobian_pinv @ delta_pose.unsqueeze(-1)
            delta_joint_pos = delta_joint_pos.squeeze(-1)
        elif self.cfg.ik_method == "svd":  # adaptive SVD
            # parameters
            k_val = self.cfg.ik_params["k_val"]
            min_singular_value = self.cfg.ik_params["min_singular_value"]
            # computation
            # U: 6xd, S: dxd, V: d x num-joint
            U, S, Vh = torch.linalg.svd(jacobian)
            S_inv = 1.0 / S
            S_inv = torch.where(S > min_singular_value, S_inv, torch.zeros_like(S_inv))
            jacobian_pinv = (
                torch.transpose(Vh, dim0=1, dim1=2)[:, :, :6]
                @ torch.diag_embed(S_inv)
                @ torch.transpose(U, dim0=1, dim1=2)
            )
            delta_joint_pos = k_val * jacobian_pinv @ delta_pose.unsqueeze(-1)
            delta_joint_pos = delta_joint_pos.squeeze(-1)
        elif self.cfg.ik_method == "trans":  # Jacobian transpose
            # parameters
            k_val = self.cfg.ik_params["k_val"]
            # computation
            jacobian_T = torch.transpose(jacobian, dim0=1, dim1=2)
            delta_joint_pos = k_val * jacobian_T @ delta_pose.unsqueeze(-1)
            delta_joint_pos = delta_joint_pos.squeeze(-1)
        elif self.cfg.ik_method == "dls":  # damped least squares
            # parameters
            lambda_val = self.cfg.ik_params["lambda_val"]
            # computation
            jacobian_T = torch.transpose(jacobian, dim0=1, dim1=2)
            lambda_matrix = (lambda_val**2) * torch.eye(n=jacobian.shape[1], device=self._device)
            delta_joint_pos = (
                jacobian_T @ torch.inverse(jacobian @ jacobian_T + lambda_matrix) @ delta_pose.unsqueeze(-1)
            )
            delta_joint_pos = delta_joint_pos.squeeze(-1)
        else:
            raise ValueError(f"Unsupported inverse-kinematics method: {self.cfg.ik_method}")

        return delta_joint_pos
//  [python]
def fit(self, dataset: Tensor) -> None:
        """Fits the PCA model to the dataset.

        Args:
          dataset (Tensor): Input dataset to fit the model.
        """
        mean = dataset.mean(dim=0)
        dataset -= mean

        _, sig, v_h = torch.svd(dataset.double())
        num_components: int
        if self.n_components <= 1:
            variance_ratios = torch.cumsum(sig * sig, dim=0) / torch.sum(sig * sig)
            num_components = torch.nonzero(variance_ratios >= self.n_components)[0]
        else:
            num_components = int(self.n_components)

        self.num_components = Tensor([num_components])

        self.singular_vectors = v_h.transpose(-2, -1)[:, :num_components].float()
        self.singular_values = sig[:num_components].float()
        self.mean = mean
//  [python]
def fit(self, dataset: Tensor) -> None:
        """Fits the PCA model to the dataset.

        Args:
          dataset (Tensor): Input dataset to fit the model.
        """
        mean = dataset.mean(dim=0)
        dataset -= mean

        _, sig, v_h = torch.linalg.svd(dataset.double(), full_matrices=False)
        num_components: int
        if self.n_components <= 1:
            variance_ratios = torch.cumsum(sig * sig, dim=0) / torch.sum(sig * sig)
            num_components = torch.nonzero(variance_ratios >= self.n_components)[0]
        else:
            num_components = int(self.n_components)

        self.num_components = Tensor([num_components])

        self.singular_vectors = v_h.transpose(-2, -1)[:, :num_components].float()
        self.singular_values = sig[:num_components].float()
        self.mean = mean
//  [python]
def transform(self, pred_coords, pred_start, pred_end, fix_start, fix_end):
        '''
        pred_coords: [batch, N, 3]
        pred_start: [batch, 3]
        pred_end: [batch, 3]
        fix_start: [batch, 3]
        fix_end: [batch, 3]
        '''
        batch = pred_coords.shape[0]
        # (batch, n, 3) 方向向量
        unit_vec = lambda x: x / torch.norm(x, dim=-1, keepdim=True)
        direct1 = unit_vec(pred_end-pred_start)
        direct2 = unit_vec(fix_end-fix_start)
        H = torch.einsum("bnc,bnd->bcd", direct1, direct2)
        # find rotation
        U, S, Vt = torch.svd(H)
        d = torch.sign(torch.det(Vt.permute(0,2,1) @ U.permute(0,2,1)))
        SS = torch.diag_embed(torch.stack([torch.ones_like(d), torch.ones_like(d), d], dim=1))
        R = (Vt.permute(0,2,1) @ SS) @ U.permute(0,2,1)

        pred_center = (pred_start + pred_end)/2
        translation = (fix_start + fix_end)/2
        pred_coords = pred_coords - pred_center.mean(dim=1,keepdim=True)

        pred_coords = (R @ pred_coords.permute(0,2,1)).permute(0,2,1) + translation.mean(dim=1,keepdim=True)
        return pred_coords
//  [python]
def transform(self, pred_coords, pred_start, pred_end, fix_start, fix_end):
        '''
        pred_coords: [batch, N, 3]
        pred_start: [batch, 3]
        pred_end: [batch, 3]
        fix_start: [batch, 3]
        fix_end: [batch, 3]
        '''
        batch = pred_coords.shape[0]
        # (batch, n, 3) 方向向量
        unit_vec = lambda x: x / torch.norm(x, dim=-1, keepdim=True)
        direct1 = unit_vec(pred_end-pred_start)
        direct2 = unit_vec(fix_end-fix_start)
        H = torch.einsum("bnc,bnd->bcd", direct1, direct2)
        # find rotation
        U, S, Vt = torch.linalg.svd(H)
        d = torch.sign(torch.linalg.det(Vt.permute(0,2,1) @ U.permute(0,2,1)))
        SS = torch.diag_embed(torch.stack([torch.ones_like(d), torch.ones_like(d), d], dim=1))
        R = (Vt.permute(0,2,1) @ SS) @ U.permute(0,2,1)

        pred_center = (pred_start + pred_end)/2
        translation = (fix_start + fix_end)/2
        pred_coords = pred_coords - pred_center.mean(dim=1,keepdim=True)

        pred_coords = (R @ pred_coords.permute(0,2,1)).permute(0,2,1) + translation.mean(dim=1,keepdim=True)
        return pred_coords
//  [python]
def opspace_matrices(mass_matrix, J_full):
    """Compute the lambda and nullspace matrices for the operational space control."""
    # Optimize code above.
    lambda_full_inv = torch.matmul(J_full, torch.linalg.solve(mass_matrix, J_full.T))

    # take the inverses, but zero out small singular values for stability
    U, S, V = torch.svd(lambda_full_inv)
    singularity_threshold = 0.05
    svd_s_inv = torch.tensor(
        [0.0 if x < singularity_threshold else float(1.0 / x) for x in S]
    ).to(mass_matrix.device)
    lambda_full = V.T.matmul(torch.diag(svd_s_inv)).matmul(U.T)

    # nullspace
    Jbar = torch.linalg.solve(mass_matrix, J_full.t()).matmul(lambda_full)
    nullspace_matrix = torch.eye(J_full.shape[-1], J_full.shape[-1]).to(
        mass_matrix.device
    ) - torch.matmul(Jbar, J_full)

    return lambda_full, nullspace_matrix
//  [python]
def opspace_matrices(mass_matrix, J_full):
    """Compute the lambda and nullspace matrices for the operational space control."""
    # Optimize code above.
    lambda_full_inv = torch.matmul(J_full, torch.linalg.solve(mass_matrix, J_full.T))

    # take the inverses, but zero out small singular values for stability
    svd_u, svd_s, svd_v = torch.linalg.svd(lambda_full_inv)
    singularity_threshold = 0.05
    svd_s_inv = torch.tensor(
        [0.0 if x < singularity_threshold else float(1.0 / x) for x in svd_s]
    ).to(mass_matrix.device)
    lambda_full = svd_v.T.matmul(torch.diag(svd_s_inv)).matmul(svd_u.T)

    # nullspace
    Jbar = torch.linalg.solve(mass_matrix, J_full.t()).matmul(lambda_full)
    nullspace_matrix = torch.eye(J_full.shape[-1], J_full.shape[-1]).to(
        mass_matrix.device
    ) - torch.matmul(Jbar, J_full)

    return lambda_full, nullspace_matrix
//  [python]
def prealign_cameras(pred, gt, verbose=False):

    if not isinstance(pred, torch.Tensor):
        pred = torch.tensor(pred).float().cuda()
    if not isinstance(gt, torch.Tensor):
        gt = torch.tensor(gt).float().cuda()

    pred = pred.clone()
    gt = gt.clone()

    T_pred = pred[:, 0:3, 3] # [N,3]
    T_gt = gt[:, 0:3, 3] # [N,3]

    R_gt = gt[:, 0:3, 0:3]
    R_pred = pred[:, 0:3, 0:3]

    try:
        sim3 = procrustes_analysis(T_gt, T_pred)
    except:
        print("warning: SVD did not converge...")
        import pdb
        pdb.set_trace()
        sim3 = edict(t0=0, t1=0, s0=1, s1=1, R=torch.eye(3).cuda())
    
    if verbose:
        print('t0 = {}, t1 = {}, s0 = {}, s1 = {}'.format(sim3.t0, sim3.t1, sim3.s0, sim3.s1))

    # align the camera poses
    T_pred = (T_pred - sim3.t1) / sim3.s1 @ sim3.R * sim3.s0 + sim3.t0 
    # draw_p(center_gt.detach().cpu(), t_aligned.detach().cpu())

    # align rotation
    U, S, Vh = torch.svd((R_pred @ R_gt.transpose(2, 1)).sum(0)) 
    R_calibrator = U @ Vh 
    if torch.det(R_calibrator) < 0:
        U[:, -1] *= -1  # Flip the sign of the last column of U
        R_calibrator = U @ Vh
    # print(rotation_distance_torch(R_calibrator @ R_pred, R_gt, degrees=True)) 
    # pose_aligned = concat_to_6D_matrix(R_calibrator @ R_pred, t_aligned[..., None])
    pred_aligned = concat_to_6D_matrix(R_pred, T_pred[..., None]) 
    return pred_aligned
//  [python]
def prealign_cameras(pred, gt, verbose=False):

    if not isinstance(pred, torch.Tensor):
        pred = torch.tensor(pred).float().cuda()
    if not isinstance(gt, torch.Tensor):
        gt = torch.tensor(gt).float().cuda()

    pred = pred.clone()
    gt = gt.clone()

    T_pred = pred[:, 0:3, 3] # [N,3]
    T_gt = gt[:, 0:3, 3] # [N,3]

    R_gt = gt[:, 0:3, 0:3]
    R_pred = pred[:, 0:3, 0:3]

    try:
        sim3 = procrustes_analysis(T_gt, T_pred)
    except:
        print("warning: SVD did not converge...")
        import pdb
        pdb.set_trace()
        sim3 = edict(t0=0, t1=0, s0=1, s1=1, R=torch.eye(3).cuda())
    
    if verbose:
        print('t0 = {}, t1 = {}, s0 = {}, s1 = {}'.format(sim3.t0, sim3.t1, sim3.s0, sim3.s1))

    # align the camera poses
    T_pred = (T_pred - sim3.t1) / sim3.s1 @ sim3.R * sim3.s0 + sim3.t0 
    # draw_p(center_gt.detach().cpu(), t_aligned.detach().cpu())

    # align rotation
    U, S, Vh = torch.linalg.svd((R_pred @ R_gt.transpose(2, 1)).sum(0)) 
    R_calibrator = U @ Vh 
    if torch.det(R_calibrator) < 0:
        U[:, -1] *= -1  # Flip the sign of the last column of U
        R_calibrator = U @ Vh
    # print(rotation_distance_torch(R_calibrator @ R_pred, R_gt, degrees=True)) 
    # pose_aligned = concat_to_6D_matrix(R_calibrator @ R_pred, t_aligned[..., None])
    pred_aligned = concat_to_6D_matrix(R_pred, T_pred[..., None]) 
    return pred_aligned
//  [python]
_IS_TORCH_BATCH_SVD_AVAILABLE = False
try:
    # Should raise an AttributeError exception if undefined.
    torch.linalg.svd
    def svd(M):
        """
        Singular Value Decomposition wrapper.
        
        Args:
            M (BxMxN tensor): batch of real matrices.
        Returns:
            (U,D,V) decomposition, such as :math:`M = U @ diag(D) @ V^T`.
        """
        U, S, V = torch.svd(M)
        return (U, S, V.transpose(-2,-1))
except (NameError, AttributeError):
    # deprecated in torch 2.0
    svd = torch.svd
//  [python]
_IS_TORCH_BATCH_SVD_AVAILABLE = False
try:
    # Should raise an AttributeError exception if undefined.
    torch.linalg.svd
    def svd(M):
        """
        Singular Value Decomposition wrapper.
        
        Args:
            M (BxMxN tensor): batch of real matrices.
        Returns:
            (U,D,V) decomposition, such as :math:`M = U @ diag(D) @ V^T`.
        """
        U, D, Vt = torch.linalg.svd(M)
        return (U, D, Vt.transpose(-2,-1))
except (NameError, AttributeError):
    # deprecated in torch 2.0
    svd = torch.svd
//  [python]
def WNNM(self,input,nsigma=100/255,c=1.41):
        U,S,Vh,= torch.svd(input)
        patch_num = self.sim_num
        C0 = c*torch.sqrt(torch.tensor([patch_num]))*2*torch.square(nsigma)
        C0 = C0.to(input.device)
        e =torch.tensor([0.00001]).to(input.device)
        Delta = torch.square(S+e)-4*C0 
        S[Delta<0] = 0
        Delta[Delta<0] = 0
        Sx = S-e+torch.sqrt(Delta)/2
        Sx[Sx<0]=0
        out = U@torch.diag_embed(Sx)@Vh
        return out
//  [python]
def WNNM(self,input,nsigma=100/255,c=1.41):
        U,S,Vh,= torch.linalg.svd(input,full_matrices=False)
        patch_num = self.sim_num
        C0 = c*torch.sqrt(torch.tensor([patch_num]))*2*torch.square(nsigma)
        C0 = C0.to(input.device)
        e =torch.tensor([0.00001]).to(input.device)
        Delta = torch.square(S+e)-4*C0 
        S[Delta<0] = 0
        Delta[Delta<0] = 0
        Sx = S-e+torch.sqrt(Delta)/2
        Sx[Sx<0]=0
        out = U@torch.diag_embed(Sx)@Vh
        return out
//  [python]
def project_so3(M, eps=1e-4):
    """
    :param M (N, *, 3, 3)
    """
    N, *dims, _, _ = M.shape
    M = M * (1 + torch.rand(N, *dims, 1, 3, device=M.device))
    U, S, Vt = torch.svd(M)  # (N, *, 3, 3), (N, *, 3), (N, *, 3, 3)
    detuvt = torch.det(torch.matmul(U, Vt))  # (N, *)
    S = torch.cat(
        [torch.ones(N, *dims, 2, device=M.device), detuvt[..., None]], dim=-1
    )  # (N, *, 3)
    return torch.matmul(U, torch.matmul(torch.diag_embed(S), Vt))
//  [python]
def project_so3(M, eps=1e-4):
    """
    :param M (N, *, 3, 3)
    """
    N, *dims, _, _ = M.shape
    M = M * (1 + torch.rand(N, *dims, 1, 3, device=M.device))
    U, D, Vt = torch.linalg.svd(M)  # (N, *, 3, 3), (N, *, 3), (N, *, 3, 3)
    detuvt = torch.linalg.det(torch.matmul(U, Vt))  # (N, *)
    S = torch.cat(
        [torch.ones(N, *dims, 2, device=M.device), detuvt[..., None]], dim=-1
    )  # (N, *, 3)
    return torch.matmul(U, torch.matmul(torch.diag_embed(S), Vt))
//  [python]
def rigid_transform_Kabsch_3D_torch(A, B):
    # assert A.shape[1] == B.shape[1]
    # num_rows, num_cols = A.shape
    # if num_rows != 3:
    #     raise Exception(f"matrix A is not 3xN, it is {num_rows}x{num_cols}")
    # num_rows, num_cols = B.shape
    # if num_rows != 3:
    #     raise Exception(f"matrix B is not 3xN, it is {num_rows}x{num_cols}")


    # find mean column wise: 3 x 1
    centroid_A = torch.mean(A, axis=-1, keepdims=True)
    centroid_B = torch.mean(B, axis=-1, keepdims=True)

    # subtract mean
    Am = A - centroid_A
    Bm = B - centroid_B

    H = Am @ Bm.transpose(-2,-1)

    # find rotation
    U, S, Vt = torch.svd(H)

    R = Vt.transpose(-2,-1) @ U.transpose(-2,-1)

    s = R.det().sign()
    R = torch.cat((R[...,:-1], (R[...,-1]*s.unsqueeze(-1)).unsqueeze(-1)), -1)

    t = -R @ centroid_A + centroid_B
    return R, t
//  [python]
def rigid_transform_Kabsch_3D_torch(A, B):
    # assert A.shape[1] == B.shape[1]
    # num_rows, num_cols = A.shape
    # if num_rows != 3:
    #     raise Exception(f"matrix A is not 3xN, it is {num_rows}x{num_cols}")
    # num_rows, num_cols = B.shape
    # if num_rows != 3:
    #     raise Exception(f"matrix B is not 3xN, it is {num_rows}x{num_cols}")


    # find mean column wise: 3 x 1
    centroid_A = torch.mean(A, axis=-1, keepdims=True)
    centroid_B = torch.mean(B, axis=-1, keepdims=True)

    # subtract mean
    Am = A - centroid_A
    Bm = B - centroid_B

    H = Am @ Bm.transpose(-2,-1)

    # find rotation
    U, S, Vt = torch.linalg.svd(H, full_matrices=False)

    R = Vt.transpose(-2,-1) @ U.transpose(-2,-1)

    s = R.det().sign()
    R = torch.cat((R[...,:-1], (R[...,-1]*s.unsqueeze(-1)).unsqueeze(-1)), -1)

    t = -R @ centroid_A + centroid_B
    return R, t
//  [python]
def get_tmscore(c1, c2):
    """
    Returns residue TM-scores for two sets of coordinate c1 and c2 in shape (n_atoms, 3)
    Directly from https://github.com/psipred/DMPfold2/blob/master/dmpfold/train.py
    """
    r1 = c1.transpose(0, 1)
    r2 = c2.transpose(0, 1)
    P = r1 - r1.mean(1).view(3, 1)
    Q = r2 - r2.mean(1).view(3, 1)
    cov = torch.matmul(P, Q.transpose(0, 1))
    try:
        U, S, Vh = torch.svd(cov)
        V = Vh.transpose(-2, -1).conj()
    except RuntimeError:
        return None
    d = torch.tensor(
        [
            [1.0, 0.0, 0.0],
            [0.0, 1.0, 0.0],
            [0.0, 0.0, torch.det(torch.matmul(V, U.transpose(0, 1)))],
        ],
        device=c1.device,
    )
    rot = torch.matmul(torch.matmul(V, d), U.transpose(0, 1))
    rot_P = torch.matmul(rot, P)
    diffs = rot_P - Q
    d0sq = ((1.24 * diffs.size(1) / 5 - 15.0) ** (1.0 / 3.0) - 1.8) ** 2
    tmscores = 1.0 / (1.0 + (diffs ** 2).sum(0) / d0sq)
    return tmscores
//  [python]
def get_tmscore(c1, c2):
    """
    Returns residue TM-scores for two sets of coordinate c1 and c2 in shape (n_atoms, 3)
    Directly from https://github.com/psipred/DMPfold2/blob/master/dmpfold/train.py
    """
    r1 = c1.transpose(0, 1)
    r2 = c2.transpose(0, 1)
    P = r1 - r1.mean(1).view(3, 1)
    Q = r2 - r2.mean(1).view(3, 1)
    cov = torch.matmul(P, Q.transpose(0, 1))
    try:
        U, _, Vh = torch.linalg.svd(cov)
        V = Vh.transpose(-2, -1).conj()
    except RuntimeError:
        return None
    d = torch.tensor(
        [
            [1.0, 0.0, 0.0],
            [0.0, 1.0, 0.0],
            [0.0, 0.0, torch.det(torch.matmul(V, U.transpose(0, 1)))],
        ],
        device=c1.device,
    )
    rot = torch.matmul(torch.matmul(V, d), U.transpose(0, 1))
    rot_P = torch.matmul(rot, P)
    diffs = rot_P - Q
    d0sq = ((1.24 * diffs.size(1) / 5 - 15.0) ** (1.0 / 3.0) - 1.8) ** 2
    tmscores = 1.0 / (1.0 + (diffs ** 2).sum(0) / d0sq)
    return tmscores
//  [python]
def svd_lowrank_approximate_tensor(x: torch.Tensor, rank=6):
    r"""This function computes a low-rank approximaiton of an input tensor on
    the last 2 dimensions

    Args:
        x: input tensor
        rank: Numerical rank to compute the approximation over

    Returns:
        Output tensor
    """
    m, n = x.shape[-2:]
    rank = min(rank, m, n)

    _x = x.reshape(-1, m, n)
    for i in range(len(_x)):
        u, s, vh = torch.svd(_x[i])
        s[rank:] = 0.0
        _x[i] = u @ torch.diag(s) @ vh
    return x
//  [python]
def svd_lowrank_approximate_tensor(x: torch.Tensor, rank=6):
    r"""This function computes a low-rank approximaiton of an input tensor on
    the last 2 dimensions

    Args:
        x: input tensor
        rank: Numerical rank to compute the approximation over

    Returns:
        Output tensor
    """
    m, n = x.shape[-2:]
    rank = min(rank, m, n)

    _x = x.reshape(-1, m, n)
    for i in range(len(_x)):
        u, s, vh = torch.linalg.svd(_x[i], full_matrices=False)
        s[rank:] = 0.0
        _x[i] = u @ torch.diag(s) @ vh
    return x
//  [python]
def solve_least_squares(A: Tensor, B: Tensor, gamma: float = 1e-6) -> Tensor:
    """
    Solves a least squares problem using SVD for stability.

    Args:
        A (Tensor): (n, m) design matrix
        B (Tensor): (n, d) right hand side

    Returns:
        Tensor: W where A @ W = B
    """
    assert A.dim() == 2, "A must be a matrix"
    assert B.dim() == 2, "B must be a matrix"
    U, S, Vh = torch.svd(A)
    return Vh.t() @ (U.t() @ B).mul((S / (S.pow(2) + gamma)).unsqueeze(-1))
//  [python]
def solve_least_squares(A: Tensor, B: Tensor, gamma: float = 1e-6) -> Tensor:
    """
    Solves a least squares problem using SVD for stability.

    Args:
        A (Tensor): (n, m) design matrix
        B (Tensor): (n, d) right hand side

    Returns:
        Tensor: W where A @ W = B
    """
    assert A.dim() == 2, "A must be a matrix"
    assert B.dim() == 2, "B must be a matrix"
    U, S, Vh = torch.linalg.svd(A, full_matrices=False)
    return Vh.t() @ (U.t() @ B).mul((S / (S.pow(2) + gamma)).unsqueeze(-1))
//  [python]
def iterate_data_rankfeat(data_loader, model, temperature):
    confs = []

    weight = model.module.classifier[1].weight.data
    B, C, H, W = weight.size()
    weight = weight.view(B, C * H * W)
    weight_sub = power_iteration(weight.unsqueeze(0), iter=100)
    weight = weight - weight_sub.squeeze()
    weight = weight.view(B, C, H, W)
    model.module.classifier[1].weight.data = weight
    
    for b, (x, y) in enumerate(data_loader):
        if b % 100 == 0:
            print('{} batches processed'.format(b))
        inputs = x.cuda()

        feat = model.module.features(inputs)
        B, C, H, W = feat.size()
        feat = feat.view(B, C, H * W)
        u, s, v = torch.svd(feat)
        feat = feat - s[:, 0:1].unsqueeze(2) * u[:, :, 0:1].bmm(v[:, 0:1, :])
        # if you want to use PI for acceleration, comment the above 2 lines and uncomment the line below
        # feat = feat - power_iteration(feat, iter=20)
        feat = feat.view(B, C, H, W)
        logits = model.module.classifier(feat)

        conf = temperature * torch.logsumexp(logits / temperature, dim=1)
        confs.extend(conf.data.cpu().numpy())

    return np.array(confs)
//  [python]
def iterate_data_rankfeat(data_loader, model, temperature):
    confs = []

    weight = model.module.classifier[1].weight.data
    B, C, H, W = weight.size()
    weight = weight.view(B, C * H * W)
    weight_sub = power_iteration(weight.unsqueeze(0), iter=100)
    weight = weight - weight_sub.squeeze()
    weight = weight.view(B, C, H, W)
    model.module.classifier[1].weight.data = weight
    
    for b, (x, y) in enumerate(data_loader):
        if b % 100 == 0:
            print('{} batches processed'.format(b))
        inputs = x.cuda()

        feat = model.module.features(inputs)
        B, C, H, W = feat.size()
        feat = feat.view(B, C, H * W)
        u, s, v = torch.linalg.svd(feat)
        feat = feat - s[:, 0:1].unsqueeze(2) * u[:, :, 0:1].bmm(v[:, 0:1, :])
        # if you want to use PI for acceleration, comment the above 2 lines and uncomment the line below
        # feat = feat - power_iteration(feat, iter=20)
        feat = feat.view(B, C, H, W)
        logits = model.module.classifier(feat)

        conf = temperature * torch.logsumexp(logits / temperature, dim=1)
        confs.extend(conf.data.cpu().numpy())

    return np.array(confs)
//  [python]
def iterate_data_rankfeat(data_loader, model, temperature):
    confs = []
    for b, (x, y) in enumerate(data_loader):
        if b % 100 == 0:
            print('{} batches processed'.format(b))
        inputs = x.cuda()

        #Block 3
        feat = model.module.features[:-2](inputs)
        B, C, H, W = feat.size()
        feat = feat.view(B, C, H * W)
        u, s, v = torch.svd(feat)
        feat = feat - s[:,0:1].unsqueeze(2)*u[:,:,0:1].bmm(v[:,0:1,:])
        # if you want to use PI for acceleration, comment the above 2 lines and uncomment the line below
        # feat = feat - power_iteration(feat, iter=20)
        feat = feat.view(B,C,H,W)
        feat = model.module.features[-2:](feat)
        logits1 = model.module.classifier(feat)

        # Block 4
        feat = model.module.features(inputs)
        B, C, H, W = feat.size()
        feat = feat.view(B, C, H * W)
        u, s, v = torch.svd(feat)
        feat = feat - s[:, 0:1].unsqueeze(2) * u[:, :, 0:1].bmm(v[:, 0:1, :])
        # if you want to use PI for acceleration, comment the above 2 lines and uncomment the line below
        # feat = feat - power_iteration(feat, iter=20)
        feat = feat.view(B, C, H, W)
        logits2 = model.module.classifier(feat)

        # Fusion at the logit space
        logits = (logits1 + logits2)/2
        conf = temperature * torch.logsumexp(logits / temperature, dim=1)
        confs.extend(conf.data.cpu().numpy())

    return np.array(confs)
//  [python]
def iterate_data_rankfeat(data_loader, model, temperature):
    confs = []
    for b, (x, y) in enumerate(data_loader):
        if b % 100 == 0:
            print('{} batches processed'.format(b))
        inputs = x.cuda()

        #Block 3
        feat = model.module.features[:-2](inputs)
        B, C, H, W = feat.size()
        feat = feat.view(B, C, H * W)
        u,s,v = torch.linalg.svd(feat)
        feat = feat - s[:,0:1].unsqueeze(2)*u[:,:,0:1].bmm(v[:,0:1,:])
        # if you want to use PI for acceleration, comment the above 2 lines and uncomment the line below
        # feat = feat - power_iteration(feat, iter=20)
        feat = feat.view(B,C,H,W)
        feat = model.module.features[-2:](feat)
        logits1 = model.module.classifier(feat)

        # Block 4
        feat = model.module.features(inputs)
        B, C, H, W = feat.size()
        feat = feat.view(B, C, H * W)
        u, s, v = torch.linalg.svd(feat)
        feat = feat - s[:, 0:1].unsqueeze(2) * u[:, :, 0:1].bmm(v[:, 0:1, :])
        # if you want to use PI for acceleration, comment the above 2 lines and uncomment the line below
        # feat = feat - power_iteration(feat, iter=20)
        feat = feat.view(B, C, H, W)
        logits2 = model.module.classifier(feat)

        # Fusion at the logit space
        logits = (logits1 + logits2)/2
        conf = temperature * torch.logsumexp(logits / temperature, dim=1)
        confs.extend(conf.data.cpu().numpy())

    return np.array(confs)
//  [python]
def svd_extract(pretrained, finetuned, rank, device="cuda", dtype=torch.float16):
    
    weight = finetuned.to(device, dtype=torch.float32) - pretrained.to(device, dtype=torch.float32)

    rank = min(rank, weight.shape[0], weight.shape[1])

    if weight.dim() == 4:
        weight = weight.reshape(weight.shape[0], -1)
    U, S, Vh = torch.svd(weight)

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

    Vh = Vh[:rank, :]

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

    U = U.clamp(low_val, hi_val)
    Vh = Vh.clamp(low_val, hi_val)

    if pretrained.dim() == 4:
        U = U.reshape(pretrained.shape[0], rank, 1, 1)
        Vh = Vh.reshape(rank, pretrained.shape[1], pretrained.shape[2], pretrained.shape[3])

    up = U.to("cpu", dtype=dtype)
    down = Vh.to("cpu", dtype=dtype)

    return up, down
//  [python]
def svd_extract(pretrained, finetuned, rank, device="cuda", dtype=torch.float16):
    
    weight = finetuned.to(device, dtype=torch.float32) - pretrained.to(device, dtype=torch.float32)

    rank = min(rank, weight.shape[0], weight.shape[1])

    if weight.dim() == 4:
        weight = weight.reshape(weight.shape[0], -1)
    U, S, Vh = torch.linalg.svd(weight)

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

    Vh = Vh[:rank, :]

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

    U = U.clamp(low_val, hi_val)
    Vh = Vh.clamp(low_val, hi_val)

    if pretrained.dim() == 4:
        U = U.reshape(pretrained.shape[0], rank, 1, 1)
        Vh = Vh.reshape(rank, pretrained.shape[1], pretrained.shape[2], pretrained.shape[3])

    up = U.to("cpu", dtype=dtype)
    down = Vh.to("cpu", dtype=dtype)

    return up, down
//  [python]
def initialize(self, V, k: int) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        'NNDSVD-A' initialization, as in sklearn.

        For nonnegative SVD initialization see https://www.cb.uu.se/~milan/histo/before2011august/Boutsidis.pdf

        NNDSVD-A first computes NNDSVD as above, then fills any zeros in W or H with the mean of V.

        :param V:
        :type V: torch.Tensor
        :param k:
        :type k: int
        :return:
        :rtype: Tuple[torch.Tensor, torch.Tensor]
        """
        U, S, Vh = torch.svd(V)

        d = U.shape[1]  # factors are identified up to sign. make sure they're positive
        for dd in range(d):
            if torch.any(U[:, dd] < 0):
                U[:, dd] *= -1
                Vh[dd, :] *= -1

        k = int(k)

        f, n = V.shape

        W = torch.zeros((f, k))
        H = torch.zeros((k, n))

        W[:, 0] = torch.sqrt(S[0]) * U[:, 0]
        H[0, :] = torch.sqrt(S[0]) * Vh[0, :]

        for j in range(1, k):
            xp = torch.nn.functional.relu(U[:, j])
            xn = torch.nn.functional.relu(-U[:, j])
            yp = torch.nn.functional.relu(Vh[j, :])
            yn = torch.nn.functional.relu(-Vh[j, :])

            mp = torch.norm(xp) * torch.norm(yp)
            mn = torch.norm(xn) * torch.norm(yn)
            if mp > mn:
                u = xp / torch.norm(xp)
                v = yp / torch.norm(yp)
                sig = mp
            else:
                u = xn / torch.norm(xn)
                v = yn / torch.norm(yn)
                sig = mn

            W[:, j] = torch.sqrt(S[j] * sig) * u
            H[j, :] = torch.sqrt(S[j] * sig) * v

        W[W == 0] = V.mean()
        H[H == 0] = V.mean()

        return W, H
//  [python]
def initialize(self, V, k: int) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        'NNDSVD-A' initialization, as in sklearn.

        For nonnegative SVD initialization see https://www.cb.uu.se/~milan/histo/before2011august/Boutsidis.pdf

        NNDSVD-A first computes NNDSVD as above, then fills any zeros in W or H with the mean of V.

        :param V:
        :type V: torch.Tensor
        :param k:
        :type k: int
        :return:
        :rtype: Tuple[torch.Tensor, torch.Tensor]
        """
        U, S, Vh = torch.linalg.svd(V, full_matrices=False)

        d = U.shape[1]  # factors are identified up to sign. make sure they're positive
        for dd in range(d):
            if torch.any(U[:, dd] < 0):
                U[:, dd] *= -1
                Vh[dd, :] *= -1

        k = int(k)

        f, n = V.shape

        W = torch.zeros((f, k))
        H = torch.zeros((k, n))

        W[:, 0] = torch.sqrt(S[0]) * U[:, 0]
        H[0, :] = torch.sqrt(S[0]) * Vh[0, :]

        for j in range(1, k):
            xp = torch.nn.functional.relu(U[:, j])
            xn = torch.nn.functional.relu(-U[:, j])
            yp = torch.nn.functional.relu(Vh[j, :])
            yn = torch.nn.functional.relu(-Vh[j, :])

            mp = torch.norm(xp) * torch.norm(yp)
            mn = torch.norm(xn) * torch.norm(yn)
            if mp > mn:
                u = xp / torch.norm(xp)
                v = yp / torch.norm(yp)
                sig = mp
            else:
                u = xn / torch.norm(xn)
                v = yn / torch.norm(yn)
                sig = mn

            W[:, j] = torch.sqrt(S[j] * sig) * u
            H[j, :] = torch.sqrt(S[j] * sig) * v

        W[W == 0] = V.mean()
        H[H == 0] = V.mean()

        return W, H
//  [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)
            U, S, V = torch.svd(Rot)
            orth = U @ V.t()
            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