jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def resize(self, if_starting, current_modes = None):
        self.current_modes = int(self.current_modes_tensor.data.item())
        self.iter = int(self.iter_tensor.data.item())
        if if_starting == True:
            matrix = torch.matmul(self.A.data[:self.current_modes+self.buffer_modes,:].T, self.B.data[:,:self.current_modes+self.buffer_modes].T)
            
            if matrix.dtype != torch.float:
                matrix = matrix.float()

            U, s, Vh = torch.linalg.svd(matrix, full_matrices=False)
            s_sqrt = torch.sqrt(s[:self.current_modes])
            del self._parameters['A']
            del self._parameters['B']
            self.A = nn.Parameter((torch.matmul(U[:,:self.current_modes], torch.diag(s_sqrt)).T).float())
            self.B = nn.Parameter((torch.matmul(torch.diag(s_sqrt), Vh[:self.current_modes,:]).T).float())

            if self.bias_flag:
                bias = self.bias.data
                del self._parameters['bias']
                self.bias = nn.Parameter(bias)
            
            del self._parameters['current_modes_tensor']
            del self._parameters['iter_tensor']
            self.current_modes_tensor = nn.Parameter(torch.tensor([self.current_modes], dtype = torch.float32))
            self.iter_tensor = nn.Parameter(torch.tensor([self.iter], dtype = torch.float32))
            self.current_modes_tensor.requires_grad = False
            self.iter_tensor.requires_grad = False

            return self.current_modes
        
        if if_starting == False:
            size_A = self.A[:current_modes,:].size()
            size_B = self.B[:,:current_modes].size()
            del self._parameters['A']
            del self._parameters['B']
            self.A = nn.Parameter(torch.empty(size_A))
            self.B = nn.Parameter(torch.empty(size_B))

            if self.bias_flag:
                bias = self.bias.data
                del self._parameters['bias']
                self.bias = nn.Parameter(bias)
            
            del self._parameters['current_modes_tensor']
            del self._parameters['iter_tensor']
            self.current_modes_tensor = nn.Parameter(torch.tensor([self.current_modes], dtype = torch.float32))
            self.iter_tensor = nn.Parameter(torch.tensor([self.iter], dtype = torch.float32))
            self.current_modes_tensor.requires_grad = False
            self.iter_tensor.requires_grad = False
//  [python]
def incremental_svd(
    A: Float[Tensor, "d r"],
    a: Float[Tensor, "batch d"],
    dtype: t.dtype = t.float64,
    n_kfac: int | None = None,
) -> Float[Tensor, "d n_kfac"]:
    """Calculate a low-rank estimate of a big [d, d] tensor, without
    materialising this full matrix.

    Args:
        A: The accumulated low-rank factor
        a: a new batch of points
        dtype: the datatype to use for the svd
        n_kfac: (optional) specify the rank of the resulting factor. If
            omitted, we use `r` from the `A` argument.
            You may choose to set this higher than the final rank during the
            accumulation.
    """
    if n_kfac is None:
        n_kfac = A.size(-1)
    a = a.to(dtype=dtype)
    A_prime = t.hstack((A, a.T))  # [d, r+batch]
    U, S, V = t.svd(A_prime, some=False)
    return U[:, :n_kfac] @ t.diag(S[:n_kfac])
//  [python]
def incremental_svd(
    A: Float[Tensor, "d r"],
    a: Float[Tensor, "batch d"],
    dtype: t.dtype = t.float64,
    n_kfac: int | None = None,
) -> Float[Tensor, "d n_kfac"]:
    """Calculate a low-rank estimate of a big [d, d] tensor, without
    materialising this full matrix.

    Args:
        A: The accumulated low-rank factor
        a: a new batch of points
        dtype: the datatype to use for the svd
        n_kfac: (optional) specify the rank of the resulting factor. If
            omitted, we use `r` from the `A` argument.
            You may choose to set this higher than the final rank during the
            accumulation.
    """
    if n_kfac is None:
        n_kfac = A.size(-1)
    a = a.to(dtype=dtype)
    A_prime = t.hstack((A, a.T))  # [d, r+batch]
    U, S, _ = t.linalg.svd(A_prime, full_matrices=False)
    return U[:, :n_kfac] @ t.diag(S[:n_kfac])
//  [python]
def svd(x, full_matrices=True, compute_uv=True):
    if not compute_uv:
        raise NotImplementedError(
            "`compute_uv=False` is not supported for torch backend."
        )
    return torch.svd(x, full_matrices=full_matrices)
//  [python]
def svd(x, full_matrices=True, compute_uv=True):
    if not compute_uv:
        raise NotImplementedError(
            "`compute_uv=False` is not supported for torch backend."
        )
    return torch.linalg.svd(x, full_matrices=full_matrices)
//  [python]
def rigid_transform_Kabsch_3D_torch(A, B):
    # R = 3x3 rotation matrix, t = 3x1 column vector
    # This already takes residue identity into account.

    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.T

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

    R = Vt.T @ U.T
    # special reflection case
    if torch.det(R) < 0:
        # print("det(R) < R, reflection detected!, correcting for it ...")
        SS = torch.diag(torch.tensor([1.,1.,-1.], device=A.device))
        R = (Vt.T @ SS) @ U.T
    assert math.fabs(torch.det(R) - 1) < 3e-3  # note I had to change this error bound to be higher

    t = -R @ centroid_A + centroid_B
    return R, t
//  [python]
def rigid_transform_Kabsch_3D_torch(A, B):
    # R = 3x3 rotation matrix, t = 3x1 column vector
    # This already takes residue identity into account.

    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.T

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

    R = Vt.T @ U.T
    # special reflection case
    if torch.linalg.det(R) < 0:
        # print("det(R) < R, reflection detected!, correcting for it ...")
        SS = torch.diag(torch.tensor([1.,1.,-1.], device=A.device))
        R = (Vt.T @ SS) @ U.T
    assert math.fabs(torch.linalg.det(R) - 1) < 3e-3  # note I had to change this error bound to be higher

    t = -R @ centroid_A + centroid_B
    return R, t
//  [python]
def weighted_triangulation(self, points2d, confi2d, R ,t):
        """
        Args:
            points2d : (V',J,2)
            confi2d : (V',J)
            R : (V',3,3)
            t : (V',3,1)
        Returns:
            points3d : (J,3)
            confi3d : (J)
        """
        n_view_filter= points2d.shape[0]
        points3d = torch.zeros((self.n_joint, 3))
        confi3d = torch.zeros((self.n_joint))
        # print(points2d.shape,confi2d.shape,R.shape,t.shape)
        for j in range(self.n_joint):
            A = []
            for i in range(n_view_filter):
                if confi2d[i,j] > 0.5:
                    P = torch.cat([R[i],t[i]],dim=1)
                    P3T = P[2]
                    A.append(confi2d[i,j] * (points2d[i,j,0]*P3T - P[0]))
                    A.append(confi2d[i,j] * (points2d[i,j,1]*P3T - P[1]))
            A = torch.stack(A)
            # print(A.shape)
            if A.shape[0] >= 4:
                u, s, vh = torch.svd(A)
                error = s[-1]
                X = vh[len(s) - 1]
                points3d[j,:] = X[:3] / X[3]
                confi3d[j] = np.exp(-torch.abs(error))
            else:
                points3d[:,j] = torch.tensor([0.0,0.0,0.0])
                confi3d[j] = 0

        return points3d, confi3d
//  [python]
def weighted_triangulation(self, points2d, confi2d, R ,t):
        """
        Args:
            points2d : (V',J,2)
            confi2d : (V',J)
            R : (V',3,3)
            t : (V',3,1)
        Returns:
            points3d : (J,3)
            confi3d : (J)
        """
        n_view_filter= points2d.shape[0]
        points3d = torch.zeros((self.n_joint, 3))
        confi3d = torch.zeros((self.n_joint))
        # print(points2d.shape,confi2d.shape,R.shape,t.shape)
        for j in range(self.n_joint):
            A = []
            for i in range(n_view_filter):
                if confi2d[i,j] > 0.5:
                    P = torch.cat([R[i],t[i]],dim=1)
                    P3T = P[2]
                    A.append(confi2d[i,j] * (points2d[i,j,0]*P3T - P[0]))
                    A.append(confi2d[i,j] * (points2d[i,j,1]*P3T - P[1]))
            A = torch.stack(A)
            # print(A.shape)
            if A.shape[0] >= 4:
                u, s, vh = torch.linalg.svd(A)
                error = s[-1]
                X = vh[len(s) - 1]
                points3d[j,:] = X[:3] / X[3]
                confi3d[j] = np.exp(-torch.abs(error))
            else:
                points3d[:,j] = torch.tensor([0.0,0.0,0.0])
                confi3d[j] = 0

        return points3d, confi3d
//  [python]
def get_orthonormal_bases_svd(vs: torch.Tensor):
    """
    Implements the solution for the Orthogonal Procrustes problem,
    which projects a matrix to the closest rotation matrix / reflection matrix using SVD.
    Args:
        vs: Tensor of shape (B, M, 9)
    Returns:
        p: Tensor of shape (B, M, 9).
    """
    # Compute SVDs of matrices in batch
    b, m, _ = vs.shape
    vs_ = vs.reshape(b * m, 3, 3)
    U, S, Vh = torch.svd(vs_)
    # Determine the diagonal matrix to make determinants 1
    sigma = torch.eye(3)[None, ...].repeat(b * m, 1, 1).to(vs_.device)
    det = torch.linalg.det(torch.bmm(U, Vh))  # Compute determinants of UVT
    ####
    # Do not set the sign of determinants to 1.
    # Inputs contain reflection matrices.
    # sigma[:, 2, 2] = det
    ####
    # Construct orthogonal matrices
    p = torch.bmm(torch.bmm(U, sigma), Vh)
    return p.reshape(b, m, 9)
//  [python]
def get_orthonormal_bases_svd(vs: torch.Tensor):
    """
    Implements the solution for the Orthogonal Procrustes problem,
    which projects a matrix to the closest rotation matrix / reflection matrix using SVD.
    Args:
        vs: Tensor of shape (B, M, 9)
    Returns:
        p: Tensor of shape (B, M, 9).
    """
    # Compute SVDs of matrices in batch
    b, m, _ = vs.shape
    vs_ = vs.reshape(b * m, 3, 3)
    U, _, Vh = torch.linalg.svd(vs_)
    # Determine the diagonal matrix to make determinants 1
    sigma = torch.eye(3)[None, ...].repeat(b * m, 1, 1).to(vs_.device)
    det = torch.linalg.det(torch.bmm(U, Vh))  # Compute determinants of UVT
    ####
    # Do not set the sign of determinants to 1.
    # Inputs contain reflection matrices.
    # sigma[:, 2, 2] = det
    ####
    # Construct orthogonal matrices
    p = torch.bmm(torch.bmm(U, sigma), Vh)
    return p.reshape(b, m, 9)
//  [python]
def rigid_transform_Kabsch_3D_torch_batch(A, B):
    # R = Bx3x3 rotation matrix, t = Bx3x1 column vector

    assert A.shape == B.shape
    _, N, M = A.shape
    if M != 3:
        raise Exception(f"matrix A and B should be BxNx3")

    A, B = A.permute(0, 2, 1), B.permute(0, 2, 1)

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

    # subtract mean
    Am = A - centroid_A
    Bm = B - centroid_B
    H = torch.bmm(Am, Bm.transpose(1, 2))

    # find rotation
    U, S, Vt = torch.svd(H)
    R = torch.bmm(Vt.transpose(1, 2), U.transpose(1, 2))

    # reflection case
    SS = torch.diag(torch.tensor([1., 1., -1.], device=A.device))
    Rm = torch.bmm(Vt.transpose(1,2) @ SS, U.transpose(1, 2))
    R = torch.where(torch.linalg.det(R)[:, None, None] < 0, Rm, R)
    assert torch.all(torch.abs(torch.linalg.det(R) - 1) < 3e-3)  # note I had to change this error bound to be higher

    t = torch.bmm(-R, centroid_A) + centroid_B
    return R, t
//  [python]
def rigid_transform_Kabsch_3D_torch_batch(A, B):
    # R = Bx3x3 rotation matrix, t = Bx3x1 column vector

    assert A.shape == B.shape
    _, N, M = A.shape
    if M != 3:
        raise Exception(f"matrix A and B should be BxNx3")

    A, B = A.permute(0, 2, 1), B.permute(0, 2, 1)

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

    # subtract mean
    Am = A - centroid_A
    Bm = B - centroid_B
    H = torch.bmm(Am, Bm.transpose(1, 2))

    # find rotation
    U, S, Vt = torch.linalg.svd(H)
    R = torch.bmm(Vt.transpose(1, 2), U.transpose(1, 2))

    # reflection case
    SS = torch.diag(torch.tensor([1., 1., -1.], device=A.device))
    Rm = torch.bmm(Vt.transpose(1,2) @ SS, U.transpose(1, 2))
    R = torch.where(torch.linalg.det(R)[:, None, None] < 0, Rm, R)
    assert torch.all(torch.abs(torch.linalg.det(R) - 1) < 3e-3)  # note I had to change this error bound to be higher

    t = torch.bmm(-R, centroid_A) + centroid_B
    return R, t
//  [python]
def calculate_erank(x):
    """
    x size: [n_nodes, dim]
    return: average erank in a batch
    """
    _, S, _ = torch.svd(x)
    N = torch.linalg.norm(S, ord=1, keepdim=True)
    S = S / N
    erank = torch.exp(torch.sum(-S * torch.log(S)))
    return erank.item()
//  [python]
def calculate_erank(x):
    """
    x size: [n_nodes, dim]
    return: average erank in a batch
    """
    _, S, _ = torch.linalg.svd(x)
    N = torch.linalg.norm(S, ord=1, keepdim=True)
    S = S / N
    erank = torch.exp(torch.sum(-S * torch.log(S)))
    return erank.item()
//  [python]
def rigid_transform_Kabsch_independent_torch(A, B):
    # R = 3x3 rotation matrix, t = 3x1 column vector
    # This already takes residue identity into account.

    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.T

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

    R = Vt.T @ U.T
    # special reflection case
    if torch.det(R) < 0:
        # print("det(R) < R, reflection detected!, correcting for it ...")
        SS = torch.diag(torch.tensor([1.,1.,-1.], device=A.device))
        R = (Vt.T @ SS) @ U.T
    assert math.fabs(torch.det(R) - 1) < 3e-3  # note I had to change this error bound to be higher

    t = - centroid_A + centroid_B # note does not change rotation
    R_vec = matrix_to_axis_angle(R)
    return t, R_vec
//  [python]
def rigid_transform_Kabsch_independent_torch(A, B):
    # R = 3x3 rotation matrix, t = 3x1 column vector
    # This already takes residue identity into account.

    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.T

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

    R = Vt.T @ U.T
    # special reflection case
    if torch.linalg.det(R) < 0:
        # print("det(R) < R, reflection detected!, correcting for it ...")
        SS = torch.diag(torch.tensor([1.,1.,-1.], device=A.device))
        R = (Vt.T @ SS) @ U.T
    assert math.fabs(torch.linalg.det(R) - 1) < 3e-3  # note I had to change this error bound to be higher

    t = - centroid_A + centroid_B # note does not change rotation
    R_vec = matrix_to_axis_angle(R)
    return t, R_vec
//  [python]
def pca(x: torch.Tensor, k: int, center=True):
    if center:
        m = x.mean(0, keepdim=True)
        s = x.std(0, unbiased=False, keepdim=True)
        x -= m
        x /= s
    U, S, V = torch.svd(x)
    reduced = torch.mm(x, V[:k].T)

    return reduced
//  [python]
def pca(x: torch.Tensor, k: int, center=True):
    if center:
        m = x.mean(0, keepdim=True)
        s = x.std(0, unbiased=False, keepdim=True)
        x -= m
        x /= s
    # why pca related to svd? https://www.cs.cmu.edu/~elaw/papers/pca.pdf chap VI
    U, S, V = torch.linalg.svd(x)
    reduced = torch.mm(x, V[:k].T)

    return reduced
//  [python]
def kabsch(coord_pred, coord_tar):
    A = coord_pred.transpose(0, 1) @ coord_tar  # [3, 3]
    U, S, Vt = torch.svd(A)
    corr_mat = torch.diag(torch.tensor([1, 1, torch.sign(torch.det(A))], device=coord_pred.device))  # [3, 3]
    rotation = (U @ corr_mat) @ Vt
    return rotation
//  [python]
def kabsch(coord_pred, coord_tar):
    A = coord_pred.transpose(0, 1) @ coord_tar  # [3, 3]
    U, S, Vt = torch.linalg.svd(A)
    corr_mat = torch.diag(torch.tensor([1, 1, torch.sign(torch.det(A))], device=coord_pred.device))  # [3, 3]
    rotation = (U @ corr_mat) @ Vt
    return rotation
//  [python]
def kabsch_batch(coords_pred, coords_tar):
    """Batch version of Kabsch algorithm."""
    A = torch.einsum("...ki, ...kj -> ...ij", coords_pred, coords_tar)
    U, S, Vt = torch.svd(A)
    sign_detA = torch.sign(torch.det(A))  # [batch_size]
    corr_mat_diag = torch.ones((A.size(0), U.size(-1)), device=A.device)  # [batch_size, 3]
    corr_mat_diag[:, -1] = sign_detA  # [batch_size, 3]
    corr_mat = torch.diag_embed(corr_mat_diag)  # [batch_size, 3, 3]
    rotation = torch.einsum("...ij, ...jk, ...kl -> ...il", U, corr_mat, Vt)  # [batch_size, 3, 3]

    return rotation
//  [python]
def kabsch_batch(coords_pred, coords_tar):
    """Batch version of Kabsch algorithm."""
    A = torch.einsum("...ki, ...kj -> ...ij", coords_pred, coords_tar)
    U, S, Vt = torch.linalg.svd(A)
    sign_detA = torch.sign(torch.det(A))  # [batch_size]
    corr_mat_diag = torch.ones((A.size(0), U.size(-1)), device=A.device)  # [batch_size, 3]
    corr_mat_diag[:, -1] = sign_detA  # [batch_size, 3]
    corr_mat = torch.diag_embed(corr_mat_diag)  # [batch_size, 3, 3]
    rotation = torch.einsum("...ij, ...jk, ...kl -> ...il", U, corr_mat, Vt)  # [batch_size, 3, 3]

    return rotation
//  [python]
def forward(self, input):
        if self.in_channels == 1:
            ba = input
        else:
            ba = BinarizeF().apply(input)

        if self.training:
            w = self.weight
            a, b = self.a, self.b
            X = w.view(w.shape[0], a, b).detach()
            if self.epoch > -1:
                for _ in range(3):
                    V = self.R1.t() @ X.detach() @ self.R2
                    B = torch.sign(V)
                    D1 = sum([Bi @ (self.R2.t()) @ (Xi.t()) for (Bi, Xi) in zip(B, X.detach())])
                    U1, S1, V1 = torch.svd(D1)
                    V1 = V1.mH
                    self.R1 = (V1 @ (U1.t()))
                    D2 = sum([(Xi.t()) @ self.R1 @ Bi for (Xi, Bi) in zip(X.detach(), B)])
                    U2, S2, V2 = torch.svd(D2)
                    V2 = V2.mH
                    self.R2 = (U2 @ (V2.t()))
            self.Rweight = ((self.R1.t()) @ X @ (self.R2)).view_as(w)
            delta = self.Rweight.detach() - w
            self.w3 = w + torch.abs(torch.sin(self.rotate)) * delta
            bw = BinarizeF().apply(self.w3)
        else:
            bw = BinarizeF().apply(self.w3)

        out = F.conv1d(input=ba, weight=bw, bias=None, stride=self.stride,
                       padding=self.padding, dilation=self.dilation, groups=self.groups)

        return out
//  [python]
def forward(self, input):
        # a = input
        # w = self.weight
        # -------IR-Net--------
        # ba = IR_Estimator().apply(a)
        # bw = IR_Estimator().apply(w)
        # -------STE----------
        # if self.training:
        if self.in_channels == 1:
            ba = input
        else:
            ba = BinarizeF().apply(input)

        # a0 = input
        if self.training:
            w = self.weight
            a, b = self.a, self.b
            X = w.view(w.shape[0], a, b).detach()
            if self.epoch > -1:
                for _ in range(3):
                    # * update B
                    V = self.R1.t() @ X.detach() @ self.R2
                    B = torch.sign(V)
                    # * update R1
                    D1 = sum([Bi @ (self.R2.t()) @ (Xi.t()) for (Bi, Xi) in zip(B, X.detach())])
                    U1, S1, V1h = torch.linalg.svd(D1)
                    V1 = V1h.mH
                    self.R1 = (V1 @ (U1.t()))
                    # * update R2
                    D2 = sum([(Xi.t()) @ self.R1 @ Bi for (Xi, Bi) in zip(X.detach(), B)])
                    U2, S2, V2h = torch.linalg.svd(D2)
                    V2 = V2h.mH
                    self.R2 = (U2 @ (V2.t()))
            self.Rweight = ((self.R1.t()) @ X @ (self.R2)).view_as(w)
            delta = self.Rweight.detach() - w
            self.w3 = w + torch.abs(torch.sin(self.rotate)) * delta
            bw = BinarizeF().apply(self.w3)
        else:
            bw = BinarizeF().apply(self.w3)
        # else:
        #     ba = torch.sign(a)
        #     bw = torch.sign(w)
        # -------2/coshx----------
        # ba = bcEstimator().apply(a)
        # bw = bcEstimator().apply(w)
        # -----------tanhx-----------
        # ba = irEstimator().apply(a)
        # bw = irEstimator().apply(w)
        # -----------fft-----------
        # ba = fft_Estimator().apply(a)
        # bw = fft_Estimator().apply(w)

        out = F.conv1d(input=ba, weight=bw, bias=None, stride=self.stride,
                       padding=self.padding, dilation=self.dilation, groups=self.groups)

        return out
//  [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))
        x = x.cuda()
        with torch.no_grad():
            # Reshape and permute the input tensor
            x = model._process_input(x)
            n = x.shape[0]

            # Expand the class token to the full batch
            batch_class_token = model.class_token.expand(n, -1, -1)
            x = torch.cat([batch_class_token, x], dim=1)

            x = x + model.encoder.pos_embedding
            x = model.encoder.dropout(x)
            
            block_lens = 12
            svd_block_index = 11
            for i in range(svd_block_index):
                x = model.encoder.layers[i](x)
            feat = x

            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_plus(feat, iter=20)

            for i in range(svd_block_index, block_lens):
                feat = model.encoder.layers[i](feat)
            feat = model.encoder.ln(feat)
        
            feat = feat[:, 0]
            logits = model.heads(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))
        x = x.cuda()
        with torch.no_grad():
            # Reshape and permute the input tensor
            x = model._process_input(x)
            n = x.shape[0]

            # Expand the class token to the full batch
            batch_class_token = model.class_token.expand(n, -1, -1)
            x = torch.cat([batch_class_token, x], dim=1)

            x = x + model.encoder.pos_embedding
            x = model.encoder.dropout(x)
            
            block_lens = 12
            svd_block_index = 11
            for i in range(svd_block_index):
                x = model.encoder.layers[i](x)
            feat = x

            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_plus(feat, iter=20)

            for i in range(svd_block_index, block_lens):
                feat = model.encoder.layers[i](feat)
            feat = model.encoder.ln(feat)
        
            feat = feat[:, 0]
            logits = model.heads(feat)

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

    return np.array(confs)
//  [python]
def forward(M : Tensor, fun : Callable[[Tensor], Tensor], fun_param : Tensor = None,
                ensure_symmetric : bool = False, ensure_psd : bool = False, mode='svd') -> Tensor:
        """Modifies the eigenvalues of a batch of symmetric matrices in the tensor M (last two dimensions).

        Source: Brooks et al. 2019, Riemannian batch normalization for SPD neural networks, NeurIPS

        Parameters
        ----------
        M : torch.Tensor
            (batch) of symmetric matrices
        fun : Callable[[Tensor], Tensor]
            elementwise function
        ensure_symmetric : bool = False (optional)
            if ensure_symmetric=True, then M is symmetrized
        ensure_psd : bool = False (optional)
            if ensure_psd=True, then the eigenvalues are clamped so that they are > 0
        -------
        Returns : torch.Tensor with modified eigenvalues
        """
        if ensure_symmetric:
            M = ensure_sym(M)

        # compute the eigenvalues and vectors
        if mode=='svd':
            U, s, _ = torch.svd(M)
        elif mode=='eigh':
            s, U = torch.linalg.eigh(M)

        if ensure_psd:
            s = s.clamp(min=EPS[s.dtype])

        # modify the eigenvalues
        smod = fun(s, fun_param)
        X = U @ torch.diag_embed(smod) @ U.transpose(-1,-2)

        return X, s, smod, U
//  [python]
def forward(M : Tensor, fun : Callable[[Tensor], Tensor], fun_param : Tensor = None,
                ensure_symmetric : bool = False, ensure_psd : bool = False, mode='svd') -> Tensor:
        """Modifies the eigenvalues of a batch of symmetric matrices in the tensor M (last two dimensions).

        Source: Brooks et al. 2019, Riemannian batch normalization for SPD neural networks, NeurIPS

        Parameters
        ----------
        M : torch.Tensor
            (batch) of symmetric matrices
        fun : Callable[[Tensor], Tensor]
            elementwise function
        ensure_symmetric : bool = False (optional)
            if ensure_symmetric=True, then M is symmetrized
        ensure_psd : bool = False (optional)
            if ensure_psd=True, then the eigenvalues are clamped so that they are > 0
        -------
        Returns : torch.Tensor with modified eigenvalues
        """
        if ensure_symmetric:
            M = ensure_sym(M)

        # compute the eigenvalues and vectors
        if mode=='svd':
            U, s, _ = torch.linalg.svd(M)
        elif mode=='eigh':
            s, U = torch.linalg.eigh(M)

        if ensure_psd:
            s = s.clamp(min=EPS[s.dtype])

        # modify the eigenvalues
        smod = fun(s, fun_param)
        X = U @ torch.diag_embed(smod) @ U.transpose(-1,-2)

        return X, s, smod, U
//  [python]
def iterate_data_rankfeatweight(data_loader, model, temperature):
    confs = []
    weight = model.encoder.layers[10].mlp[3].weight.data
    B, C = weight.size()
    weight_sub = power_iteration_plus(weight.unsqueeze(0), iter=100)
    weight = weight - weight_sub.squeeze()
    model.encoder.layers[10].mlp[3].weight.data = weight
    
    for b, (x, y) in enumerate(data_loader):
        if b % 100 == 0:
            print('{} batches processed'.format(b))
        x = x.cuda()

        with torch.no_grad():
            # Reshape and permute the input tensor
            x = model._process_input(x)
            n = x.shape[0]

            # Expand the class token to the full batch
            batch_class_token = model.class_token.expand(n, -1, -1)
            x = torch.cat([batch_class_token, x], dim=1)


            x = x + model.encoder.pos_embedding
            x = model.encoder.dropout(x)
            
            block_lens = 12
            svd_block_index = 11
            for i in range(svd_block_index):
                x = model.encoder.layers[i](x)
            feat = x

            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_plus(feat, iter=20)

            for i in range(svd_block_index, block_lens):
                feat = model.encoder.layers[i](feat)
            feat = model.encoder.ln(feat)
        
            feat = feat[:, 0]
            logits = model.heads(feat)

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

    return np.array(confs)
//  [python]
def iterate_data_rankfeatweight(data_loader, model, temperature):
    confs = []
    weight = model.encoder.layers[10].mlp[3].weight.data
    B, C = weight.size()
    weight_sub = power_iteration_plus(weight.unsqueeze(0), iter=100)
    weight = weight - weight_sub.squeeze()
    model.encoder.layers[10].mlp[3].weight.data = weight
    
    for b, (x, y) in enumerate(data_loader):
        if b % 100 == 0:
            print('{} batches processed'.format(b))
        x = x.cuda()

        with torch.no_grad():
            # Reshape and permute the input tensor
            x = model._process_input(x)
            n = x.shape[0]

            # Expand the class token to the full batch
            batch_class_token = model.class_token.expand(n, -1, -1)
            x = torch.cat([batch_class_token, x], dim=1)


            x = x + model.encoder.pos_embedding
            x = model.encoder.dropout(x)
            
            block_lens = 12
            svd_block_index = 11
            for i in range(svd_block_index):
                x = model.encoder.layers[i](x)
            feat = x

            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_plus(feat, iter=20)

            for i in range(svd_block_index, block_lens):
                feat = model.encoder.layers[i](feat)
            feat = model.encoder.ln(feat)
        
            feat = feat[:, 0]
            logits = model.heads(feat)

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

    return np.array(confs)
//  [python]
def forward(self, protein_coords_pred: List[Tensor], protein_coords: List[Tensor]) -> Tensor:
        rmsds = []
        for protein_coords_pred, protein_coords in zip(protein_coords_pred, protein_coords):
            protein_coords_pred_mean = protein_coords_pred.mean(dim=0, keepdim=True)  # (1,3)
            protein_coords_mean = protein_coords.mean(dim=0, keepdim=True)  # (1,3)

            A = (protein_coords_pred - protein_coords_pred_mean).transpose(0, 1) @ (
                        protein_coords - protein_coords_mean)

            U, S, Vt = torch.svd(A)

            # corr_mat = torch.diag(torch.tensor([1, 1, torch.sign(torch.det(A))], device=protein_coords_pred.device))
            corr_mat = torch.diag(
                torch.tensor([1, 1, torch.sign(torch.det(Vt.t() @ U.t()))], device=protein_coords_pred.device))
            rotation = (U @ corr_mat) @ Vt
            # corr_mat = torch.diag(torch.tensor([1, 1, torch.sign(torch.det(Vt.t() @ U.t()))], device=protein_coords_pred.device))
            # rotation = (Vt.t() @ corr_mat) @ U.t()
            translation = protein_coords_pred_mean - torch.t(rotation @ protein_coords_mean.t())  # (1,3)

            protein_coords = (rotation @ protein_coords.t()).t() + translation

            rmsds.append(torch.sqrt(torch.mean(torch.sum(((protein_coords_pred - protein_coords) ** 2), dim=1))))
        return torch.tensor(rmsds).mean()
//  [python]
def forward(self, protein_coords_pred: List[Tensor], protein_coords: List[Tensor]) -> Tensor:
        rmsds = []
        for protein_coords_pred, protein_coords in zip(protein_coords_pred, protein_coords):
            protein_coords_pred_mean = protein_coords_pred.mean(dim=0, keepdim=True)  # (1,3)
            protein_coords_mean = protein_coords.mean(dim=0, keepdim=True)  # (1,3)

            A = (protein_coords_pred - protein_coords_pred_mean).transpose(0, 1) @ (
                        protein_coords - protein_coords_mean)

            U, S, Vt = torch.linalg.svd(A)

            # corr_mat = torch.diag(torch.tensor([1, 1, torch.sign(torch.det(A))], device=protein_coords_pred.device))
            corr_mat = torch.diag(
                torch.tensor([1, 1, torch.sign(torch.det(Vt.t() @ U.t()))], device=protein_coords_pred.device))
            rotation = (U @ corr_mat) @ Vt
            # corr_mat = torch.diag(torch.tensor([1, 1, torch.sign(torch.det(Vt.t() @ U.t()))], device=protein_coords_pred.device))
            # rotation = (Vt.t() @ corr_mat) @ U.t()
            translation = protein_coords_pred_mean - torch.t(rotation @ protein_coords_mean.t())  # (1,3)

            protein_coords = (rotation @ protein_coords.t()).t() + translation

            rmsds.append(torch.sqrt(torch.mean(torch.sum(((protein_coords_pred - protein_coords) ** 2), dim=1))))
        return torch.tensor(rmsds).mean()
//  [python]
def get_matrix(self, p1, p2, w=None):
        """
        Find R and T which is the solution to argmin_{R, T} \sum_i ||p2_i - (R * p1_i + T)||_2
        See https://ieeexplore.ieee.org/document/4767965


        Args:
          x, y: [n_batch, n_points, dim]
          w: [n_batch, n_points]
        Returns:
          A: [n_batch, dim, dim+1]
        """
        # Take transpose as columns should be the points
        p1 = p1.permute(0, 2, 1)
        p2 = p2.permute(0, 2, 1)

        # Calculate centroids
        if w is not None:
            # find weighed mean column wise
            p1_c = torch.sum(p1 * w, axis=2, keepdims=True)
            p2_c = torch.sum(p2 * w, axis=2, keepdims=True)
        else:
            p1_c = torch.mean(p1, axis=2, keepdim=True)
            p2_c = torch.mean(p2, axis=2, keepdim=True)

        # Subtract centroids
        q1 = p1 - p1_c
        q2 = p2 - p2_c

        if w is not None:
            q1 = q1 * w
            q2 = q2 * w

        # Calculate covariance matrix
        H = torch.bmm(q1, q2.transpose(1, 2))

        # Calculate singular value decomposition (SVD)
        U, S, V = torch.svd(H)  # the SVD of linalg gives you Vt
        Ut = U.transpose(1, 2)
        Vt = V.transpose(1, 2)

        R = torch.bmm(Vt, Ut)

        # assert torch.allclose(
        #     torch.linalg.det(R), torch.tensor(1.0)
        # ), f"Rotation matrix of N-point registration not 1, see paper Arun et al., det={torch.linalg.det(R)}"

        # special reflection case
        dets = torch.det(R)
        dets = torch.unsqueeze(dets, -1)
        dets = torch.stack([torch.ones_like(dets), torch.ones_like(dets), dets], axis=1)
        dets = torch.cat([dets, dets, dets], axis=2)
        Vt = Vt * torch.sign(dets)

        # Calculate rotation matrix
        R = torch.bmm(Vt, Ut)

        # Calculate translation matrix
        T = p2_c - torch.bmm(R, p1_c)

        # Create augmented matrix
        aug_mat = torch.cat([R, T], axis=-1)
        return aug_mat
//  [python]
def get_matrix(self, p1, p2, w=None):
        """
        Find R and T which is the solution to argmin_{R, T} \sum_i ||p2_i - (R * p1_i + T)||_2
        See https://ieeexplore.ieee.org/document/4767965


        Args:
          x, y: [n_batch, n_points, dim]
          w: [n_batch, n_points]
        Returns:
          A: [n_batch, dim, dim+1]
        """
        # Take transpose as columns should be the points
        p1 = p1.permute(0, 2, 1)
        p2 = p2.permute(0, 2, 1)

        # Calculate centroids
        if w is not None:
            # find weighed mean column wise
            p1_c = torch.sum(p1 * w, axis=2, keepdims=True)
            p2_c = torch.sum(p2 * w, axis=2, keepdims=True)
        else:
            p1_c = torch.mean(p1, axis=2, keepdim=True)
            p2_c = torch.mean(p2, axis=2, keepdim=True)

        # Subtract centroids
        q1 = p1 - p1_c
        q2 = p2 - p2_c

        if w is not None:
            q1 = q1 * w
            q2 = q2 * w

        # Calculate covariance matrix
        H = torch.bmm(q1, q2.transpose(1, 2))

        # Calculate singular value decomposition (SVD)
        U, _, Vt = torch.linalg.svd(H)  # the SVD of linalg gives you Vt
        Ut = U.transpose(1, 2)
        V = Vt.transpose(1, 2)
        R = torch.bmm(V, Ut)

        # assert torch.allclose(
        #     torch.linalg.det(R), torch.tensor(1.0)
        # ), f"Rotation matrix of N-point registration not 1, see paper Arun et al., det={torch.linalg.det(R)}"

        # special reflection case
        dets = torch.det(R)
        dets = torch.unsqueeze(dets, -1)
        dets = torch.stack([torch.ones_like(dets), torch.ones_like(dets), dets], axis=1)
        dets = torch.cat([dets, dets, dets], axis=2)
        V = V * torch.sign(dets)

        # Calculate rotation matrix
        R = torch.bmm(V, Ut)

        # Calculate translation matrix
        T = p2_c - torch.bmm(R, p1_c)

        # Create augmented matrix
        aug_mat = torch.cat([R, T], axis=-1)
        return aug_mat
//  [python]
def low_rank_decomposition(
        weight: torch.Tensor,
        rank_ratio: Optional[float] = 0.1,
        parameter_ratio: Optional[float] = 0.15,
        remove_criteria: Optional[str] = 'max_eigenvalue',
        return_dict: Optional[bool] = False,
):
    """
    Parameters
    ----------
    weight: torch.Tensor
        The matrix to decompose, of shape (H, W)
    rank_ratio: float, optional, default 0.1
        The ratio of the reduced rank to the original rank:
            rank_of_decomposed_matrix / rank_of_input_weight
    parameter_ratio: float, optional, default 0.15
        The ratio of the number of parameters of the decomposed matrix to the original matrix:
            parameter_num_of_decomposed_matrix / (H * W).
        If specify, override rank_ratio
    remove_criteria: str, optional, default 'max_eigenvalue'
        The criteria to remove the small eigenvalues, of ['max_eigenvalue', 'random', 'min_eigenvalue']
    return_dict: bool, optional, default False
        Return a dict if True, else return a tuple (L, R)
    debug: bool, optional, default False
        Print debug information if True
    """
    matrix_dimension = len(weight.size())
    assert matrix_dimension == 2, "Only Support 2D matrix"
    height, width = weight.size()

    # Use SVD to decompose a matrix, default full_matrices is False to save parameters
    dtype = weight.dtype
    weight = weight.float()
    u_, s_, v_ = torch.svd(weight)
    u_ = u_.to(dtype)
    s_ = s_.to(dtype)
    v_ = v_.to(dtype)
    rank = torch.count_nonzero(s_)

    if parameter_ratio is not None:
        reduced_rank = math.ceil(parameter_ratio * (height * width) / (height + width))
    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, :]) @ v_
    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, :]) @ v_
    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:, :]) @ v_
    else:
        raise NameError("remove criteria not support")

    if return_dict:
        return {"L": l_, "R": r_, "U": u_, "S": s_, "Vh": v_, 'reduced_rank': reduced_rank}
    else:
        return l_, r_
//  [python]
def low_rank_decomposition(
        weight: torch.Tensor,
        rank_ratio: Optional[float] = 0.1,
        parameter_ratio: Optional[float] = 0.15,
        remove_criteria: Optional[str] = 'max_eigenvalue',
        return_dict: Optional[bool] = False,
):
    """
    Parameters
    ----------
    weight: torch.Tensor
        The matrix to decompose, of shape (H, W)
    rank_ratio: float, optional, default 0.1
        The ratio of the reduced rank to the original rank:
            rank_of_decomposed_matrix / rank_of_input_weight
    parameter_ratio: float, optional, default 0.15
        The ratio of the number of parameters of the decomposed matrix to the original matrix:
            parameter_num_of_decomposed_matrix / (H * W).
        If specify, override rank_ratio
    remove_criteria: str, optional, default 'max_eigenvalue'
        The criteria to remove the small eigenvalues, of ['max_eigenvalue', 'random', 'min_eigenvalue']
    return_dict: bool, optional, default False
        Return a dict if True, else return a tuple (L, R)
    debug: bool, optional, default False
        Print debug information if True
    """
    matrix_dimension = len(weight.size())
    assert matrix_dimension == 2, "Only Support 2D matrix"
    height, width = weight.size()

    # Use SVD to decompose a matrix, default full_matrices is False to save parameters
    dtype = weight.dtype
    weight = weight.float()
    u_, s_, v_ = torch.linalg.svd(weight, full_matrices=False)
    u_ = u_.to(dtype)
    s_ = s_.to(dtype)
    v_ = v_.to(dtype)
    rank = torch.count_nonzero(s_)

    if parameter_ratio is not None:
        reduced_rank = math.ceil(parameter_ratio * (height * width) / (height + width))
    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, :]) @ v_
    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, :]) @ v_
    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:, :]) @ v_
    else:
        raise NameError("remove criteria not support")

    if return_dict:
        return {"L": l_, "R": r_, "U": u_, "S": s_, "Vh": v_, 'reduced_rank': reduced_rank}
    else:
        return l_, r_
//  [python]
def lplr_svd(
    X: torch.Tensor = None,
    r: int = None,
    B1: int = 8,
    B2: int = 8,
    normalize_and_shift=False,
    seed=None,
    sketch: Sketch = Sketch.GAUSSIAN,
    **kwargs,
) -> torch.Tensor:
    """
    :param X: Target matrix
    :param r: Output rank
    :param B1: Bit-budget for first low-rank factor
    :param B2: Bit-budget for second low-rank factor
    :param q: No. of power iterations
    :param sketch: Specify the sketching matrix use for the first low-rank factor
    :return: Low-precision Low-rank approximation
    """

    if seed is not None:
        logger.trace(f"Using seed = {seed}")
        torch.random.manual_seed(seed)

    orig_dtype = X.dtype
    X = X.float()  # Do simulations in float

    # Compute SVD
    U, S, V = torch.svd(X.float())

    # Extract singular values and singular vectors
    U = U[:, 0:r]
    Σ = torch.diag(S[0:r])
    VT = V[:, 0:r].t()

    sketch = Sketch(sketch)
    # Compute the first quantized low-rank factor
    match sketch:
        case Sketch.GAUSSIAN:
            Z = U @ Σ

            # Sketch the column space of Z with S matrix
            S = torch.randn(Z.shape[1], r, device=X.device, dtype=X.dtype) / math.sqrt(
                r
            )  # Gaussian sketching matrix

            Z = quantize(Z @ S, B=B1, preserve_original_dtype=True)
        case Sketch.NOTHING:
            Z = U @ Σ
            Z = quantize(Z, B=B1, preserve_original_dtype=True)
        case Sketch.SPARSE_JL:
            Z = U @ Σ
            S = (
                sparse_jl_transform(Z.shape[1], r, kwargs.get("sparse_jl_s"))
                .to(Z.device)
                .type(Z.dtype)
            )
            Z = quantize(Z @ S, B=B1, preserve_original_dtype=True)
        case _:
            logger.error(f"Unsupported sketch type {sketch}")

    # Compute the second low-rank factor
    pinv = torch.linalg.pinv(Z.float()).type(X.dtype)
    if torch.isnan(pinv).any().item():
        logger.error(f"NaNs encountered in pinv")

    W = pinv @ X

    if torch.isnan(W).any().item():
        logger.error(f"NaNs encountered in pinv @ X")

    # Quantize the second low-rank factor
    W = quantize(W, B=B2, preserve_original_dtype=True)

    out = Z @ W
    if normalize_and_shift:
        out = normalize_and_shift_wrt_inner_prod(X, Z @ W)

    if torch.isnan(out).any().item():
        logger.error(f"NaNs encountered in LPLRed matrix")
    out = out.type(orig_dtype)

    return out
//  [python]
def lplr_svd(
    X: torch.Tensor = None,
    r: int = None,
    B1: int = 8,
    B2: int = 8,
    normalize_and_shift=False,
    seed=None,
    sketch: Sketch = Sketch.GAUSSIAN,
    **kwargs,
) -> torch.Tensor:
    """
    :param X: Target matrix
    :param r: Output rank
    :param B1: Bit-budget for first low-rank factor
    :param B2: Bit-budget for second low-rank factor
    :param q: No. of power iterations
    :param sketch: Specify the sketching matrix use for the first low-rank factor
    :return: Low-precision Low-rank approximation
    """

    if seed is not None:
        logger.trace(f"Using seed = {seed}")
        torch.random.manual_seed(seed)

    orig_dtype = X.dtype
    X = X.float()  # Do simulations in float

    # Compute SVD
    U, Σ, VT = torch.linalg.svd(X.float(), full_matrices=False)

    # Extract singular values and singular vectors
    U = U[:, 0:r]
    Σ = Σ[0:r]
    VT = VT[0:r, :]

    sketch = Sketch(sketch)
    # Compute the first quantized low-rank factor
    match sketch:
        case Sketch.GAUSSIAN:
            Z = U @ torch.diag(Σ)

            # Sketch the column space of Z with S matrix
            S = torch.randn(Z.shape[1], r, device=X.device, dtype=X.dtype) / math.sqrt(
                r
            )  # Gaussian sketching matrix

            Z = quantize(Z @ S, B=B1, preserve_original_dtype=True)
        case Sketch.NOTHING:
            Z = U @ torch.diag(Σ)
            Z = quantize(Z, B=B1, preserve_original_dtype=True)
        case Sketch.SPARSE_JL:
            Z = U @ torch.diag(Σ)
            S = (
                sparse_jl_transform(Z.shape[1], r, kwargs.get("sparse_jl_s"))
                .to(Z.device)
                .type(Z.dtype)
            )
            Z = quantize(Z @ S, B=B1, preserve_original_dtype=True)
        case _:
            logger.error(f"Unsupported sketch type {sketch}")

    # Compute the second low-rank factor
    pinv = torch.linalg.pinv(Z.float()).type(X.dtype)
    if torch.isnan(pinv).any().item():
        logger.error(f"NaNs encountered in pinv")

    W = pinv @ X

    if torch.isnan(W).any().item():
        logger.error(f"NaNs encountered in pinv @ X")

    # Quantize the second low-rank factor
    W = quantize(W, B=B2, preserve_original_dtype=True)

    out = Z @ W
    if normalize_and_shift:
        out = normalize_and_shift_wrt_inner_prod(X, Z @ W)

    if torch.isnan(out).any().item():
        logger.error(f"NaNs encountered in LPLRed matrix")
    out = out.type(orig_dtype)

    return out