jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def uniform_init_(tensor):
    """Fills in the input ``tensor`` in place with an orthogonal matrix.
    If square, the matrix will have positive determinant.
    The tensor will be distributed according to the Haar measure.
    The input tensor must have at least 2 dimensions.
    For tensors with more than 2 dimensions the first dimensions are treated as
    batch dimensions.

    Args:
        tensor (torch.Tensor): a 2-dimensional tensor or a batch of them
    """
    if tensor.ndim < 2:
        raise ValueError('Only tensors with 2 or more dimensions are supported. Got a tensor of shape {}'.format(tuple(tensor.size())))
    n, k = tensor.size()[-2:]
    transpose = n < k
    with torch.no_grad():
        x = torch.empty_like(tensor).normal_(0, 1)
        if transpose:
            x.transpose_(-2, -1)
        q, r = torch.qr(x)
        d = r.diagonal(dim1=-2, dim2=-1).sign()
        q *= d.unsqueeze(-2)
        if transpose:
            q.transpose_(-2, -1)
        if n == k:
            mask = (torch.det(q) >= 0.0).float()
            mask[mask == 0.0] = -1.0
            mask = mask.unsqueeze(-1)
            q[..., 0] *= mask
        tensor.copy_(q)
        return tensor
//  [python]
def uniform_init_(tensor):
    """Fills in the input ``tensor`` in place with an orthogonal matrix.
    If square, the matrix will have positive determinant.
    The tensor will be distributed according to the Haar measure.
    The input tensor must have at least 2 dimensions.
    For tensors with more than 2 dimensions the first dimensions are treated as
    batch dimensions.

    Args:
        tensor (torch.Tensor): a 2-dimensional tensor or a batch of them
    """
    if tensor.ndim < 2:
        raise ValueError('Only tensors with 2 or more dimensions are supported. Got a tensor of shape {}'.format(tuple(tensor.size())))
    n, k = tensor.size()[-2:]
    transpose = n < k
    with torch.no_grad():
        x = torch.empty_like(tensor).normal_(0, 1)
        if transpose:
            x.transpose_(-2, -1)
        q, r = torch.linalg.qr(x)
        d = r.diagonal(dim1=-2, dim2=-1).sign()
        q *= d.unsqueeze(-2)
        if transpose:
            q.transpose_(-2, -1)
        if n == k:
            mask = (torch.det(q) >= 0.0).float()
            mask[mask == 0.0] = -1.0
            mask = mask.unsqueeze(-1)
            q[..., 0] *= mask
        tensor.copy_(q)
        return tensor
//  [python]
def tangent_basis(normal_vec):
    A = torch.eye(normal_vec.shape[0]).to(normal_vec.device, normal_vec.dtype)
    A[:, 0] = normal_vec
    Q, R = torch.qr(A)
    norms = torch.linalg.norm(normal_vec)
    Q[:, 0] = normal_vec / torch.linalg.norm(normal_vec)
    return Q.nan_to_num()
//  [python]
def tangent_basis(normal_vec):
    A = torch.eye(normal_vec.shape[0]).to(normal_vec.device, normal_vec.dtype)
    A[:, 0] = normal_vec
    Q, R = torch.linalg.qr(A)
    norms = torch.linalg.norm(normal_vec)
    Q[:, 0] = normal_vec / torch.linalg.norm(normal_vec)
    return Q.nan_to_num()
//  [python]
def samples_from_haar_distribution(d, bsize, device, dtype) -> torch.Tensor:
    """Random unitary matrices drawn from U(d) Haar distribution
    Adopted from scipy.stats.unitary_group, which implements the algorithm described in
    Mezzadri, How to generate random matrices from the classical compact groups (2006)
    """
    z = 1 / math.sqrt(2) * (
        torch.randn(bsize, d, d, device=device, dtype=dtype) +
        1j * torch.randn(bsize, d, d, device=device, dtype=dtype)
    )
    q, r = torch.qr(z)
    # The last two dimensions are the rows and columns of R matrices.
    # Extract the diagonals. Note that this eliminates a dimension.
    d = r.diagonal(offset=0, dim1=-2, dim2=-1)
    # Add back a dimension for proper broadcasting: we're dividing
    # each row of each R matrix by the diagonal of the R matrix.
    q *= (d/abs(d))[..., None, :]  # to broadcast properly
    return q
//  [python]
def samples_from_haar_distribution(d, bsize, device, dtype) -> torch.Tensor:
    """Random unitary matrices drawn from U(d) Haar distribution
    Adopted from scipy.stats.unitary_group, which implements the algorithm described in
    Mezzadri, How to generate random matrices from the classical compact groups (2006)
    """
    z = 1 / math.sqrt(2) * (
        torch.randn(bsize, d, d, device=device, dtype=dtype) +
        1j * torch.randn(bsize, d, d, device=device, dtype=dtype)
    )
    q, r = torch.linalg.qr(z)
    # The last two dimensions are the rows and columns of R matrices.
    # Extract the diagonals. Note that this eliminates a dimension.
    d = r.diagonal(offset=0, dim1=-2, dim2=-1)
    # Add back a dimension for proper broadcasting: we're dividing
    # each row of each R matrix by the diagonal of the R matrix.
    q *= (d/abs(d))[..., None, :]  # to broadcast properly
    return q
//  [python]
def bminres(A, B, M=None, X0=None, tol=1e-5, maxiters=None, verbose=True):
    """
    O'Leary (1980) Algorithm B-CG Version 4 (Paige and Saunders form)
    """
    m, n = A.shape
    n2, s = B.shape
    assert m == n
    assert n2 == n
    #if M is None:
    #    M = lo.IdentityOperator(m)
    if maxiters is None:
        maxiters = n
    if X0 is None:
        X0 = torch.zeros_like(B)

    k = 1

    R_hat_0 = B - A @ X0

    X_hat = X0
    Z_km1 = torch.zeros_like(X0)
    Z_k, nu_k_inv = torch.qr(R_hat_0)
    RHS_psi = nu_k_inv
    psi_km1 = torch.zeros(s, s)
    W_bar_k = Z_k
    #rho_km1 = 0
    AZ_k = A @ Z_k
    rho_k = Z_k.T @ AZ_k
    L_bar_k_k = rho_k
    V_k_T = torch.eye(2 * s)

    for k in range(1, maxiters):
        Z_k, nu_k_inv, Z_km1 = *torch.qr(AZ_k - Z_k @ rho_k - Z_km1 @ nu_k_inv.T), Z_k
        print(torch.linalg.norm(torch.eye(s) - Z_k.T @ Z_k))
        AZ_k = A @ Z_k
        rho_k = Z_k.T @ AZ_k
        temp = nu_k_inv @ V_k_T[s:, :] # Double check this step
        L_k_km2, L_bar_kp1_k = temp[:, :s], temp[:, s:]

        temp = torch.vstack([L_bar_k_k.T, nu_k_inv])
        V_k_T, L_k_k_T_aug = torch.qr(temp, 'complete') # V_k_T should not be orthogonal. V_k should be.
        L_k_k = L_k_k_T_aug[:s, :].T
        print(torch.linalg.norm(torch.eye(2 * s) - V_k_T @ V_k_T.T))
        print(torch.linalg.norm(torch.eye(2 * s) - V_k_T.T @ V_k_T))

        temp = torch.hstack([L_bar_kp1_k, rho_k]) @ V_k_T
        L_kp1_k, L_bar_k_k = temp[:, s:], temp[:, :s]


        temp = torch.hstack([W_bar_k, Z_k]) @ V_k_T
        W_k, W_bar_k = temp[:, s:], temp[:, :s]

        # TODO: Check if better solve; should be triangular
        psi_k = torch.linalg.solve(L_k_k, RHS_psi)
        RHS_psi = -(L_kp1_k @ psi_k + L_k_km2 @ psi_km1)

        X_hat += W_k @ psi_k
        psi_km1 = psi_k
        del psi_k

        psi_bar_kp1 = torch.linalg.solve(L_bar_k_k, RHS_psi)
        X = X_hat + W_bar_k  @ psi_bar_kp1
        print(k, torch.linalg.norm(B - A @ X) / torch.linalg.norm(B))

    # TODO: Check if better solve; should be triangular
    psi_bar_kp1 = torch.linalg.solve(L_bar_k_k, RHS_psi)
    return X_hat + W_bar_k  @ psi_bar_kp1, k + 1
//  [python]
def bminres(A, B, M=None, X0=None, tol=1e-5, maxiters=None, verbose=True):
    """
    O'Leary (1980) Algorithm B-CG Version 4 (Paige and Saunders form)
    """
    m, n = A.shape
    n2, s = B.shape
    assert m == n
    assert n2 == n
    #if M is None:
    #    M = lo.IdentityOperator(m)
    if maxiters is None:
        maxiters = n
    if X0 is None:
        X0 = torch.zeros_like(B)

    k = 1

    R_hat_0 = B - A @ X0

    X_hat = X0
    Z_km1 = torch.zeros_like(X0)
    Z_k, nu_k_inv = torch.linalg.qr(R_hat_0)
    RHS_psi = nu_k_inv
    psi_km1 = torch.zeros(s, s)
    W_bar_k = Z_k
    #rho_km1 = 0
    AZ_k = A @ Z_k
    rho_k = Z_k.T @ AZ_k
    L_bar_k_k = rho_k
    V_k_T = torch.eye(2 * s)

    for k in range(1, maxiters):
        Z_k, nu_k_inv, Z_km1 = *torch.linalg.qr(AZ_k - Z_k @ rho_k - Z_km1 @ nu_k_inv.T), Z_k
        print(torch.linalg.norm(torch.eye(s) - Z_k.T @ Z_k))
        AZ_k = A @ Z_k
        rho_k = Z_k.T @ AZ_k
        temp = nu_k_inv @ V_k_T[s:, :] # Double check this step
        L_k_km2, L_bar_kp1_k = temp[:, :s], temp[:, s:]

        temp = torch.vstack([L_bar_k_k.T, nu_k_inv])
        V_k_T, L_k_k_T_aug = torch.linalg.qr(temp, 'complete') # V_k_T should not be orthogonal. V_k should be.
        L_k_k = L_k_k_T_aug[:s, :].T
        print(torch.linalg.norm(torch.eye(2 * s) - V_k_T @ V_k_T.T))
        print(torch.linalg.norm(torch.eye(2 * s) - V_k_T.T @ V_k_T))

        temp = torch.hstack([L_bar_kp1_k, rho_k]) @ V_k_T
        L_kp1_k, L_bar_k_k = temp[:, s:], temp[:, :s]


        temp = torch.hstack([W_bar_k, Z_k]) @ V_k_T
        W_k, W_bar_k = temp[:, s:], temp[:, :s]

        # TODO: Check if better solve; should be triangular
        psi_k = torch.linalg.solve(L_k_k, RHS_psi)
        RHS_psi = -(L_kp1_k @ psi_k + L_k_km2 @ psi_km1)

        X_hat += W_k @ psi_k
        psi_km1 = psi_k
        del psi_k

        psi_bar_kp1 = torch.linalg.solve(L_bar_k_k, RHS_psi)
        X = X_hat + W_bar_k  @ psi_bar_kp1
        print(k, torch.linalg.norm(B - A @ X) / torch.linalg.norm(B))

    # TODO: Check if better solve; should be triangular
    psi_bar_kp1 = torch.linalg.solve(L_bar_k_k, RHS_psi)
    return X_hat + W_bar_k  @ psi_bar_kp1, k + 1
//  [python]
def random_rotation_matrix(batch_size: int, ndim=3, generator=None) -> torch.Tensor:
    z = torch.randn((batch_size, ndim, ndim), generator=generator)
    q, r = torch.qr(z)
    sign = 2 * (r.diagonal(dim1=-2, dim2=-1) >= 0) - 1
    rot = q
    rot *= sign[..., None, :]
    rot[:, 0, :] *= torch.det(rot)[..., None]
    return rot
//  [python]
def random_rotation_matrix(batch_size: int, ndim=3, generator=None) -> torch.Tensor:
    z = torch.randn((batch_size, ndim, ndim), generator=generator)
    q, r = torch.linalg.qr(z)
    sign = 2 * (r.diagonal(dim1=-2, dim2=-1) >= 0) - 1
    rot = q
    rot *= sign[..., None, :]
    rot[:, 0, :] *= torch.linalg.det(rot)[..., None]
    return rot
//  [python]
def forward(self, input):
        # input: (bs*t) x 512
        weight = self.weight + 1e-8
        Q, R = torch.qr(weight)  # get eignvector, orthogonal [n1, n2, n3, n4]

        if input is None:
            return Q
        else:
            input_diag = torch.diag_embed(input)  # alpha, diagonal matrix
            out = torch.matmul(input_diag, Q.T)
            out = torch.sum(out, dim=1)

            return out
//  [python]
def forward(self, input):
        # input: (bs*t) x 512
        weight = self.weight + 1e-8
        Q, R = torch.linalg.qr(weight)  # get eignvector, orthogonal [n1, n2, n3, n4]

        if input is None:
            return Q
        else:
            input_diag = torch.diag_embed(input)  # alpha, diagonal matrix
            out = torch.matmul(input_diag, Q.T)
            out = torch.sum(out, dim=1)

            return out
//  [python]
def forward(ctx, A, b, cache_decomposition=True):
        B, M, N = A.shape
        assert b.shape == (B, M, 1)

        with torch.no_grad():
            Q, R = torch.qr(A, some=True)
            x = torch.linalg.solve_triangular(R, torch.bmm(b.view(B, 1, M), Q).view(B, N, 1), upper=True)

        # save state for backward pass
        ctx.save_for_backward(A, b, x, R if cache_decomposition else None)

        # return solution
        return x
//  [python]
def forward(ctx, A, b, cache_decomposition=True):
        B, M, N = A.shape
        assert b.shape == (B, M, 1)

        with torch.no_grad():
            Q, R = torch.linalg.qr(A, mode='reduced')
            x = torch.linalg.solve_triangular(R, torch.bmm(b.view(B, 1, M), Q).view(B, N, 1), upper=True)

        # save state for backward pass
        ctx.save_for_backward(A, b, x, R if cache_decomposition else None)

        # return solution
        return x
//  [python]
def _sample_symmetric(dim, mean=0, std=10, **kwargs):
    """generate a random symmetric matrix"""
    # sample random orthogonal matrix (uniformly)
    Q, R = torch.qr(torch.randn(dim, dim, **kwargs))
    Q.mul_(R.diagonal().sign().view(1, -1))

    # sample eigenvalues from normal distribution
    e = torch.empty(dim, **kwargs).normal_(mean, std)

    return Q @ torch.diag(e) @ Q.T
//  [python]
def _sample_symmetric(dim, mean=0, std=10, **kwargs):
    """generate a random symmetric matrix"""
    # sample random orthogonal matrix (uniformly)
    Q, R = torch.linalg.qr(torch.randn(dim, dim, **kwargs))
    Q.mul_(R.diagonal().sign().view(1, -1))

    # sample eigenvalues from normal distribution
    e = torch.empty(dim, **kwargs).normal_(mean, std)

    return Q @ torch.diag(e) @ Q.T
//  [python]
def test_fisher_vec_prod(self, state):
        x, y, x_dim, y_dim, k = state

        V, _ = torch.qr(torch.ones((x_dim, k), device=DEVICE))
        U = V.clone()
        Jv = autodiff.jacobian_vector_product(y, x, V)
        Fv = autodiff.vector_jacobian_product(y, x, Jv)

        jac = autodiff.jacobian(y, x)

        Fv2 = jac.T @ jac @ U  # manually compute product to compare accuracy

        assert Fv.shape == (x_dim, k)
        assert Fv2.allclose(Fv, atol=1E-6)
//  [python]
def test_fisher_vec_prod(self, state):
        x, y, x_dim, y_dim, k = state

        V, _ = torch.linalg.qr(torch.ones((x_dim, k), device=DEVICE), "reduced")
        U = V.clone()
        Jv = autodiff.jacobian_vector_product(y, x, V)
        Fv = autodiff.vector_jacobian_product(y, x, Jv)

        jac = autodiff.jacobian(y, x)

        Fv2 = jac.T @ jac @ U  # manually compute product to compare accuracy

        assert Fv.shape == (x_dim, k)
        assert Fv2.allclose(Fv, atol=1E-6)
//  [python]
def _solve_preconditioner(self) -> Optional[Callable]:
        r"""
        (Optional) define a preconditioner :math:`\mathbf P` that can be used for linear systems,
        but not necessarily for log determinants. By default, this can call
        :meth:`~linear_operator.LinearOperator._preconditioner`.

        :return: a function which performs :math:`\mathbf P^{-1}(\cdot)`
        """
        base_precond, _, _ = self._preconditioner()

        if base_precond is not None:
            return base_precond
        elif linear_operator.beta_features.default_preconditioner.on():
            if hasattr(self, "_default_preconditioner_cache"):
                U, S, Vt = self._default_preconditioner_cache
            else:
                precond_basis_size = min(linear_operator.settings.max_preconditioner_size.value(), self.size(-1))
                random_basis = torch.randn(
                    self.batch_shape + torch.Size((self.size(-2), precond_basis_size)),
                    device=self.device,
                    dtype=self.dtype,
                )
                projected_mat = self._matmul(random_basis)
                proj_q, _ = torch.qr(projected_mat)
                orthog_projected_mat = self._matmul(proj_q).mT
                # Maybe log
                if settings.verbose_linalg.on():
                    settings.verbose_linalg.logger.debug(
                        f"Running svd on a matrix of size {orthog_projected_mat.shape}."
                    )
                U, S, Vt = torch.linalg.svd(orthog_projected_mat)
                U = proj_q.matmul(U)

                self._default_preconditioner_cache = (U, S, Vt)

            def preconditioner(v):
                res = Vt.matmul(v)
                res = (1 / S).unsqueeze(-1) * res
                res = U.matmul(res)
                return res

            return preconditioner
        else:
            return None
//  [python]
def _solve_preconditioner(self) -> Optional[Callable]:
        r"""
        (Optional) define a preconditioner :math:`\mathbf P` that can be used for linear systems,
        but not necessarily for log determinants. By default, this can call
        :meth:`~linear_operator.LinearOperator._preconditioner`.

        :return: a function which performs :math:`\mathbf P^{-1}(\cdot)`
        """
        base_precond, _, _ = self._preconditioner()

        if base_precond is not None:
            return base_precond
        elif linear_operator.beta_features.default_preconditioner.on():
            if hasattr(self, "_default_preconditioner_cache"):
                U, S, Vt = self._default_preconditioner_cache
            else:
                precond_basis_size = min(linear_operator.settings.max_preconditioner_size.value(), self.size(-1))
                random_basis = torch.randn(
                    self.batch_shape + torch.Size((self.size(-2), precond_basis_size)),
                    device=self.device,
                    dtype=self.dtype,
                )
                projected_mat = self._matmul(random_basis)
                proj_q = torch.linalg.qr(projected_mat)
                orthog_projected_mat = self._matmul(proj_q).mT
                # Maybe log
                if settings.verbose_linalg.on():
                    settings.verbose_linalg.logger.debug(
                        f"Running svd on a matrix of size {orthog_projected_mat.shape}."
                    )
                U, S, Vt = torch.linalg.svd(orthog_projected_mat)
                U = proj_q.matmul(U)

                self._default_preconditioner_cache = (U, S, Vt)

            def preconditioner(v):
                res = Vt.matmul(v)
                res = (1 / S).unsqueeze(-1) * res
                res = U.matmul(res)
                return res

            return preconditioner
        else:
            return None
//  [python]
def get_residue_coordinate_systems(ca_positions, c_positions, n_positions):
    ca_c_vec = c_positions - ca_positions  # B x 3
    ca_n_vec = n_positions - ca_positions  # B x 3
    q, r = torch.linalg.qr(
        torch.stack((ca_c_vec, ca_n_vec), dim=1).transpose(2, 1)
    )
    orth_basis = q.transpose(2, 1)

    # Basis 1 is normalized ca_o_vec and basis 2 is the orthogonal projection of basis 1 on basis 2
    basis_1 = orth_basis[:, 0]
    basis_2 = orth_basis[:, 1]
    basis_3 = torch.cross(basis_1, basis_2)

    basis_matrix = torch.stack((basis_1, basis_2, basis_3), dim=2)
    inv_basis_matrix = torch.linalg.inv(basis_matrix)
    return ResidueCoordinateSystem(
        basis_matrix=basis_matrix, inv_basis_matrix=inv_basis_matrix
    )
//  [python]
def get_residue_coordinate_systems(ca_positions, c_positions, n_positions):
    ca_c_vec = c_positions - ca_positions  # B x 3
    ca_n_vec = n_positions - ca_positions  # B x 3
    orth_basis = torch.linalg.qr(
        torch.stack((ca_c_vec, ca_n_vec), dim=1).transpose(2, 1)
    )[0]
    orth_basis = orth_basis.transpose(2, 1)

    # Basis 1 is normalized ca_o_vec and basis 2 is the orthogonal projection of basis 1 on basis 2
    basis_1 = orth_basis[:, 0]
    basis_2 = orth_basis[:, 1]
    basis_3 = torch.cross(basis_1, basis_2)

    basis_matrix = torch.stack((basis_1, basis_2, basis_3), dim=2)
    inv_basis_matrix = torch.linalg.inv(basis_matrix)
    return ResidueCoordinateSystem(
        basis_matrix=basis_matrix, inv_basis_matrix=inv_basis_matrix
    )
//  [python]
def _orthogonalize(matrices, epsilon=0):
    """
    Decide between Gram-Schmidt or QR factorization to orthogonalize a batch of matrices.
    QR factorization doesn't work with half-precision, but it is usually faster with a rank > 2.
    """
    assert len(matrices.shape) == 3 and matrices.shape[2] <= matrices.shape[1]

    num_matrices = matrices.shape[0]
    rank = matrices.shape[2]
    dtype = matrices.dtype
    if rank <= 2 or dtype in [torch.float16, torch.bfloat16]:
        _orthogonalize_gram_schmidt(matrices, epsilon=epsilon)
    else:
        torch.qr(
            matrices,
            out=(
                matrices,
                torch.empty(num_matrices, rank, rank, device=matrices.device, dtype=dtype)
            )
        )
//  [python]
def _orthogonalize(matrices, epsilon=0):
    """
    Decide between Gram-Schmidt or QR factorization to orthogonalize a batch of matrices.
    QR factorization doesn't work with half-precision, but it is usually faster with a rank > 2.
    """
    assert len(matrices.shape) == 3 and matrices.shape[2] <= matrices.shape[1]

    num_matrices = matrices.shape[0]
    rank = matrices.shape[2]
    dtype = matrices.dtype
    if rank <= 2 or dtype in [torch.float16, torch.bfloat16]:
        _orthogonalize_gram_schmidt(matrices, epsilon=epsilon)
    else:
        torch.linalg.qr(
            matrices,
            out=(
                matrices,
                torch.empty(num_matrices, rank, rank, device=matrices.device, dtype=dtype)
            )
        )
//  [python]
def _orthogonal_matrix_chunk(cols, device = None):
    unstructured_block = torch.randn((cols, cols), device = device)
    q, r = torch.qr(unstructured_block.cpu())
    q, r = map(lambda t: t.to(device), (q, r))
    return q.t()
//  [python]
def _orthogonal_matrix_chunk(cols, device = None):
    unstructured_block = torch.randn((cols, cols), device = device)
    q, r = torch.linalg.qr(unstructured_block.cpu(), mode = 'complete') 
    q, r = map(lambda t: t.to(device), (q, r))
    return q.t()
//  [python]
def test_qr_error_cases(self, device, dtype):
        t1 = torch.randn(5, device=device, dtype=dtype)
        with self.assertRaisesRegex(RuntimeError, 'linalg.qr: The input tensor A must have at least 2 dimensions.'):
            torch.qr(t1)
        t2 = torch.randn((5, 7), device=device, dtype=dtype)
        with self.assertRaisesRegex(RuntimeError, "qr received unrecognized mode 'hello'"):
            torch.qr(t2, mode='hello')
//  [python]
def test_qr_error_cases(self, device, dtype):
        t1 = torch.randn(5, device=device, dtype=dtype)
        with self.assertRaisesRegex(RuntimeError, 'linalg.qr: The input tensor A must have at least 2 dimensions.'):
            torch.linalg.qr(t1)
        t2 = torch.randn((5, 7), device=device, dtype=dtype)
        with self.assertRaisesRegex(RuntimeError, "qr received unrecognized mode 'hello'"):
            torch.linalg.qr(t2, mode='hello')
//  [python]
def __init__(
            self, batch, rank, n, is_thin, init_mode,
            init_std=0.0001, init_seed=None, init_qr_mul_sgn_diag_r=False,
            handle_det_sign_pos=False,
    ):
        super().__init__()
        assert 0 < rank <= n
        self.batch = batch
        self.rank = rank
        self.n = n
        self.is_thin = is_thin and (rank != n)
        self.num_param = batch * dof_stiefel_canonical(rank, n)
        rng = np.random if init_seed is None else np.random.RandomState(init_seed)
        self.param = torch.nn.Parameter(torch.zeros(batch, n-rank, rank))
        self.register_buffer('eye', torch.eye(n, rank).unsqueeze(0).repeat(batch, 1, 1))
        if handle_det_sign_pos and rank == n and n % 2 == 1:
            self.register_buffer('negate_last_col', torch.ones(1, 1, n))
            self.negate_last_col[-1] = -1
        with torch.no_grad():
            if init_mode == 'eye':
                self.param.zero_()
            elif init_mode in ('qr_eye_randn', 'qr_randn'):
                if init_mode == 'qr_eye_randn':
                    tmp = (
                        self.eye +
                        init_std * torch.from_numpy(rng.randn(batch, n, rank)).clamp(min=-1, max=1)
                    )
                else:
                    tmp = torch.from_numpy(rng.randn(batch, n, rank))
                param = []
                for m in tmp:
                    q, r = torch.qr(m)
                    if init_qr_mul_sgn_diag_r:
                        q *= r.diag().sign().view(1, rank)
                    hh = q.geqrf()[0].unsqueeze(0)
                    param.append(hh[:, rank:, :])
                param = torch.cat(param, dim=0)
                self.param.copy_(param)
            else:
                raise ValueError(f'Unknown orthogonal initialization {init_mode}')
//  [python]
def __init__(
            self, batch, rank, n, is_thin, init_mode,
            init_std=0.0001, init_seed=None, init_qr_mul_sgn_diag_r=False,
            handle_det_sign_pos=False,
    ):
        super().__init__()
        assert 0 < rank <= n
        self.batch = batch
        self.rank = rank
        self.n = n
        self.is_thin = is_thin and (rank != n)
        self.num_param = batch * dof_stiefel_canonical(rank, n)
        rng = np.random if init_seed is None else np.random.RandomState(init_seed)
        self.param = torch.nn.Parameter(torch.zeros(batch, n-rank, rank))
        self.register_buffer('eye', torch.eye(n, rank).unsqueeze(0).repeat(batch, 1, 1))
        if handle_det_sign_pos and rank == n and n % 2 == 1:
            self.register_buffer('negate_last_col', torch.ones(1, 1, n))
            self.negate_last_col[-1] = -1
        with torch.no_grad():
            if init_mode == 'eye':
                self.param.zero_()
            elif init_mode in ('qr_eye_randn', 'qr_randn'):
                if init_mode == 'qr_eye_randn':
                    tmp = (
                        self.eye +
                        init_std * torch.from_numpy(rng.randn(batch, n, rank)).clamp(min=-1, max=1)
                    )
                else:
                    tmp = torch.from_numpy(rng.randn(batch, n, rank))
                param = []
                for m in tmp:
                    q, r = torch.linalg.qr(m)
                    if init_qr_mul_sgn_diag_r:
                        q *= r.diag().sign().view(1, rank)
                    hh = q.geqrf()[0].unsqueeze(0)
                    param.append(hh[:, rank:, :])
                param = torch.cat(param, dim=0)
                self.param.copy_(param)
            else:
                raise ValueError(f'Unknown orthogonal initialization {init_mode}')
//  [python]
def __init__(
            self, batch, rank, n, is_thin, init_mode,
            init_std=0.0001, init_seed=None, init_qr_mul_sgn_diag_r=False,
            handle_det_sign_pos=False,
    ):
        super().__init__()
        assert 0 < rank <= n
        self.batch = batch
        self.rank = rank
        self.n = n
        self.is_thin = is_thin and (rank != n)
        self.num_param = batch * dof_stiefel(rank, n)
        rng = np.random if init_seed is None else np.random.RandomState(init_seed)
        self.param = torch.nn.Parameter(torch.zeros(batch, n, rank))
        self.register_buffer('eye', torch.eye(n, rank).unsqueeze(0).repeat(batch, 1, 1))
        if handle_det_sign_pos and rank == n and n % 2 == 1:
            self.register_buffer('negate_last_col', torch.ones(1, 1, n))
            self.negate_last_col[-1] = -1
        with torch.no_grad():
            if init_mode == 'eye':
                self.param.zero_()
            elif init_mode in ('qr_eye_randn', 'qr_randn'):
                if init_mode == 'qr_eye_randn':
                    tmp = (
                        self.eye +
                        init_std * torch.from_numpy(rng.randn(batch, n, rank)).clamp(min=-1, max=1)
                    )
                else:
                    tmp = torch.from_numpy(rng.randn(batch, n, rank))
                param = []
                for m in tmp:
                    q, r = torch.qr(m)
                    if init_qr_mul_sgn_diag_r:
                        q *= r.diag().sign().view(1, rank)
                    hh = q.geqrf()[0].unsqueeze(0)
                    param.append(hh)
                param = torch.cat(param, dim=0)
                self.param.copy_(param)
            else:
                raise ValueError(f'Unknown orthogonal initialization {init_mode}')
//  [python]
def __init__(
            self, batch, rank, n, is_thin, init_mode,
            init_std=0.0001, init_seed=None, init_qr_mul_sgn_diag_r=False,
            handle_det_sign_pos=False,
    ):
        super().__init__()
        assert 0 < rank <= n
        self.batch = batch
        self.rank = rank
        self.n = n
        self.is_thin = is_thin and (rank != n)
        self.num_param = batch * dof_stiefel(rank, n)
        rng = np.random if init_seed is None else np.random.RandomState(init_seed)
        self.param = torch.nn.Parameter(torch.zeros(batch, n, rank))
        self.register_buffer('eye', torch.eye(n, rank).unsqueeze(0).repeat(batch, 1, 1))
        if handle_det_sign_pos and rank == n and n % 2 == 1:
            self.register_buffer('negate_last_col', torch.ones(1, 1, n))
            self.negate_last_col[-1] = -1
        with torch.no_grad():
            if init_mode == 'eye':
                self.param.zero_()
            elif init_mode in ('qr_eye_randn', 'qr_randn'):
                if init_mode == 'qr_eye_randn':
                    tmp = (
                        self.eye +
                        init_std * torch.from_numpy(rng.randn(batch, n, rank)).clamp(min=-1, max=1)
                    )
                else:
                    tmp = torch.from_numpy(rng.randn(batch, n, rank))
                param = []
                for m in tmp:
                    q, r = torch.linalg.qr(m)
                    if init_qr_mul_sgn_diag_r:
                        q *= r.diag().sign().view(1, rank)
                    hh = q.geqrf()[0].unsqueeze(0)
                    param.append(hh)
                param = torch.cat(param, dim=0)
                self.param.copy_(param)
            else:
                raise ValueError(f'Unknown orthogonal initialization {init_mode}')
//  [python]
def test_linalg_qr():
    N0 = 5
    N1 = 3
    matA = torch.randn(N0, N1, dtype=torch.float64) + 1j*torch.randn(N0, N1, dtype=torch.float64)
    matQ,matR = torch.qr(matA, some=True)
    assert torch.abs(matA - matQ @ matR).max().item() < 1e-10
    assert torch.abs(torch.tril(matR, -1)).max() < 1e-10
    assert torch.abs(matQ.T @ matQ.conj() - torch.eye(N1, dtype=torch.float64)).max() < 1e-10

    hf0 = lambda matA,matB: torch.dot(matB.reshape(-1), torch.qr(matA, some=True)[0].reshape(-1)).real
    matB = torch.randn(N0, N1, dtype=torch.float64) + 1j*torch.randn(N0, N1, dtype=torch.float64)
    tmp0 = matA.clone()
    tmp0.requires_grad_(True)
    ret0 = torch.autograd.grad(hf0(tmp0, matB), [tmp0])[0]
    ret_ = torch.zeros_like(matA)
    zero_eps = 1e-4
    for ind0 in range(N0):
        for ind1 in range(N1):
            tmp0,tmp1,tmp2,tmp3 = [matA.clone() for _ in range(4)]
            tmp0[ind0,ind1] += zero_eps
            tmp1[ind0,ind1] -= zero_eps
            tmp2[ind0,ind1] += zero_eps*1j
            tmp3[ind0,ind1] -= zero_eps*1j
            ret_[ind0,ind1] = (hf0(tmp0, matB) - hf0(tmp1, matB) + 1j*(hf0(tmp2, matB) - hf0(tmp3, matB)))/(2*zero_eps)
    assert torch.abs(ret0 - ret_).max().item() < 1e-6
//  [python]
def test_linalg_qr():
    N0 = 5
    N1 = 3
    matA = torch.randn(N0, N1, dtype=torch.float64) + 1j*torch.randn(N0, N1, dtype=torch.float64)
    matQ,matR = torch.linalg.qr(matA, mode='reduced')
    assert torch.abs(matA - matQ @ matR).max().item() < 1e-10
    assert torch.abs(torch.tril(matR, -1)).max() < 1e-10
    assert torch.abs(matQ.T @ matQ.conj() - torch.eye(N1, dtype=torch.float64)).max() < 1e-10

    hf0 = lambda matA,matB: torch.dot(matB.reshape(-1), torch.linalg.qr(matA, mode='reduced')[0].reshape(-1)).real
    matB = torch.randn(N0, N1, dtype=torch.float64) + 1j*torch.randn(N0, N1, dtype=torch.float64)
    tmp0 = matA.clone()
    tmp0.requires_grad_(True)
    ret0 = torch.autograd.grad(hf0(tmp0, matB), [tmp0])[0]
    ret_ = torch.zeros_like(matA)
    zero_eps = 1e-4
    for ind0 in range(N0):
        for ind1 in range(N1):
            tmp0,tmp1,tmp2,tmp3 = [matA.clone() for _ in range(4)]
            tmp0[ind0,ind1] += zero_eps
            tmp1[ind0,ind1] -= zero_eps
            tmp2[ind0,ind1] += zero_eps*1j
            tmp3[ind0,ind1] -= zero_eps*1j
            ret_[ind0,ind1] = (hf0(tmp0, matB) - hf0(tmp1, matB) + 1j*(hf0(tmp2, matB) - hf0(tmp3, matB)))/(2*zero_eps)
    assert torch.abs(ret0 - ret_).max().item() < 1e-6
//  [python]
def QRreduction(datas):
    ndatas = torch.qr(datas.permute(0, 2, 1))[1]
    ndatas = ndatas.permute(0, 2, 1)
    return ndatas
//  [python]
def QRreduction(datas):
    ndatas = torch.linalg.qr(datas.permute(0, 2, 1),'reduced').R
    ndatas = ndatas.permute(0, 2, 1)
    return ndatas
//  [python]
def forward(self, x):
        with torch.cuda.amp.autocast(self.fp16):
            x = self.conv1(x)
            x = self.bn1(x)
            x = self.prelu(x)
            x = self.layer1(x)
            x = self.layer2(x)
            x = self.layer3(x)
            x = self.layer4(x)
            x = self.bn2(x)
            x = torch.flatten(x, 1)
            x = self.dropout(x)
        if self.make_orthonormal:
            q, r = torch.qr(self.ort_weight)
            x = F.linear(x, q.t())
        else:
            x = self.fc(x.float() if self.fp16 else x)
        x = self.features(x)
        return x
//  [python]
def forward(self, x):
        with torch.cuda.amp.autocast(self.fp16):
            x = self.conv1(x)
            x = self.bn1(x)
            x = self.prelu(x)
            x = self.layer1(x)
            x = self.layer2(x)
            x = self.layer3(x)
            x = self.layer4(x)
            x = self.bn2(x)
            x = torch.flatten(x, 1)
            x = self.dropout(x)
        if self.make_orthonormal:
            x = F.linear(x, torch.linalg.qr(self.ort_weight)[0].t())
        else:
            x = self.fc(x.float() if self.fp16 else x)
        x = self.features(x)
        return x
//  [python]
def image_based_lighting(surf: torch.Tensor,
                         norm: torch.Tensor,
                         view: torch.Tensor,
                         probe: dotdict,  # lighting
                         albedo: float,  # scalar albedo
                         roughness: float,  # scalar roughness
                         microfacet: Microfacet,  # material
                         N: int = 1024,  # number of samples
                         H: int = 16,
                         W: int = 32,
                         uniform: bool = True,  # uniform or stratified sampling
                         perturb: bool = True,
                         ):
    # Generate sample_count uniformly and stratified samples over the sphere
    P = surf.shape[0]
    if uniform:  # uniform random sampling
        T = P * N
        theta, phi = spherical_uniform_sampling_upper(T, device=surf.device)  # T, T,
        ray_d = spher2cart(theta, phi)  # T, 3, z always bigger than zero
    else:  # stratified sampling
        N = H * W
        T = P * N
        xyz, area = gen_light_xyz(H, W, device=surf.device)
        
        if perturb:
            R = torch.rand(3, 3, device=xyz.device)
            Q, R = torch.qr(R)  # 3, 3
            xyz = xyz @ Q  # apply random rotation
            
        xyz, area = xyz.view(-1, 3), area.view(-1, 1)
        ray_d = normalize(xyz)  # T, 3
        # Adding more samples seems to help, but not very good for low roughness surface (i.e. this implementation has a hard upper limit for specular surfaces)
        # And the visibility's influence is not clear enough, you've only done ablation on one of the char in one of the novel lighting
        # The physical correctness of distance field soft shadow is quesitonable

        # __import__('ipdb').set_trace()
        # torch.testing.assert_allclose(probe.view(-1, 3), sample_envmap_image(probe, ray_d))

        xyz = xyz[:, None].expand(N, P, 3).reshape(-1, 3)  # T, 3
        area = area[:, None].expand(N, P, 1).reshape(-1, 1)  # T, 1
        ray_d = ray_d[:, None].expand(N, P, 3).reshape(-1, 3)  # T, 3

    # Preparing shapes
    norm = norm[None].expand(N, P, 3).reshape(T, 3)  # T, 3
    view = view[None].expand(N, P, 3).reshape(T, 3)  # T, 3

    # Transform ray_d to be pointing upward from normal direction
    if uniform:
        R = torch.zeros([T, 3, 3], device=norm.device)
        R[..., 0, 0] = 1.0
        R[..., :3, 2] = norm  # c2w, z axis is normal direction
        R[..., :3, 1] = normalize(torch.cross(R[..., :3, 2], R[..., :3, 0]))
        R[..., :3, 0] = normalize(torch.cross(R[..., :3, 1], R[..., :3, 2]))
        ray_d = (R @ ray_d[..., None])[..., 0]

    # Compute shading
    ldot = (ray_d * norm).sum(dim=-1, keepdim=True)  # T
    light = sample_envmap_image(probe, ray_d)
    brdf = microfacet(ray_d, -view, norm, albedo, roughness)
    shading = light * ldot * brdf

    # Apply area to normalize integration
    if uniform:
        shading = shading * 2.0 * torch.pi / N
    else:
        shading = shading * area
    shading = shading.view(N, P, -1).sum(dim=-3)

    return shading
//  [python]
def image_based_lighting(surf: torch.Tensor,
                         norm: torch.Tensor,
                         view: torch.Tensor,
                         probe: dotdict,  # lighting
                         albedo: float,  # scalar albedo
                         roughness: float,  # scalar roughness
                         microfacet: Microfacet,  # material
                         N: int = 1024,  # number of samples
                         H: int = 16,
                         W: int = 32,
                         uniform: bool = True,  # uniform or stratified sampling
                         perturb: bool = True,
                         ):
    # Generate sample_count uniformly and stratified samples over the sphere
    P = surf.shape[0]
    if uniform:  # uniform random sampling
        T = P * N
        theta, phi = spherical_uniform_sampling_upper(T, device=surf.device)  # T, T,
        ray_d = spher2cart(theta, phi)  # T, 3, z always bigger than zero
    else:  # stratified sampling
        N = H * W
        T = P * N
        xyz, area = gen_light_xyz(H, W, device=surf.device)
        
        if perturb:
            R = torch.rand(3, 3, device=xyz.device)
            Q, R = torch.linalg.qr(R)  # 3, 3
            xyz = xyz @ Q  # apply random rotation
            
        xyz, area = xyz.view(-1, 3), area.view(-1, 1)
        ray_d = normalize(xyz)  # T, 3
        # Adding more samples seems to help, but not very good for low roughness surface (i.e. this implementation has a hard upper limit for specular surfaces)
        # And the visibility's influence is not clear enough, you've only done ablation on one of the char in one of the novel lighting
        # The physical correctness of distance field soft shadow is quesitonable

        # __import__('ipdb').set_trace()
        # torch.testing.assert_allclose(probe.view(-1, 3), sample_envmap_image(probe, ray_d))

        xyz = xyz[:, None].expand(N, P, 3).reshape(-1, 3)  # T, 3
        area = area[:, None].expand(N, P, 1).reshape(-1, 1)  # T, 1
        ray_d = ray_d[:, None].expand(N, P, 3).reshape(-1, 3)  # T, 3

    # Preparing shapes
    norm = norm[None].expand(N, P, 3).reshape(T, 3)  # T, 3
    view = view[None].expand(N, P, 3).reshape(T, 3)  # T, 3

    # Transform ray_d to be pointing upward from normal direction
    if uniform:
        R = torch.zeros([T, 3, 3], device=norm.device)
        R[..., 0, 0] = 1.0
        R[..., :3, 2] = norm  # c2w, z axis is normal direction
        R[..., :3, 1] = normalize(torch.cross(R[..., :3, 2], R[..., :3, 0]))
        R[..., :3, 0] = normalize(torch.cross(R[..., :3, 1], R[..., :3, 2]))
        ray_d = (R @ ray_d[..., None])[..., 0]

    # Compute shading
    ldot = (ray_d * norm).sum(dim=-1, keepdim=True)  # T
    light = sample_envmap_image(probe, ray_d)
    brdf = microfacet(ray_d, -view, norm, albedo, roughness)
    shading = light * ldot * brdf

    # Apply area to normalize integration
    if uniform:
        shading = shading * 2.0 * torch.pi / N
    else:
        shading = shading * area
    shading = shading.view(N, P, -1).sum(dim=-3)

    return shading
//  [python]
def lr_torch_new(x1, x2):
    x1_flat, x2_flat = rearrange_activations(x1), rearrange_activations(x2)
    q2, r2 = torch.qr(x2_flat)
    return (torch.linalg.norm(q2.T @ x1_flat)) ** 2 / (torch.linalg.norm(x1_flat)) ** 2
//  [python]
def lr_torch_new(x1, x2):
    x1_flat, x2_flat = rearrange_activations(x1), rearrange_activations(x2)
    q2, r2 = torch.linalg.qr(x2_flat)
    return (torch.linalg.norm(q2.T @ x1_flat)) ** 2 / (torch.linalg.norm(x1_flat)) ** 2
//  [python]
def sum_kronecker_products(
        left_tensor: Tensor,
        right_tensor: Tensor,
        assert_positive_definite: bool = True) -> Tuple[Tensor, Tensor]:
    r"""Approximation of sums of general Kronecker products with the power method.

    It approximates
        .. math:
            L, R = \text{arg min}_{L, R} \| \frac{1}{B} \sum_{i = 1}^B
                left_tensor[i] \otimes right_tensor[i] - L \otimes R\|

    It corresponds to a generalized version of Algorithm 2 of the appendix of
    `Natural continual learning: success is a journey, not (just) a destination
    <https://proceedings.neurips.cc/paper/2021/file/ec5aa0b7846082a2415f0902f0da88f2-Paper.pdf>_`

    Args:
        left_tensor: (shape [B, M, M]) The first tensor
        right_tensor: (shape [B, N, N]) The second tensor
        assert_positive_definite: If the result should be made positive definite
            to counteract numeric instabilities

    Returns:
        The two tensors L (shape [M, M]) and R (shape [N, N])
    """
    l = left_tensor.view(left_tensor.shape[0], -1)
    r = right_tensor.view(right_tensor.shape[0], -1)
    Q, _ = torch.qr(l.T)
    H = (l @ Q).T @ r
    U, s, VT = torch.linalg.svd(H, full_matrices=False)
    y = s[0].sqrt() * VT[0]
    x = s[0].sqrt() * Q @ U[:, 0]
    L, R = x.view(*left_tensor.shape[1:]), y.view(*right_tensor.shape[1:])
    if torch.trace(L) < 0 and torch.trace(R) < 0:
        L, R = -L, -R
    if assert_positive_definite:
        L, R = check_and_make_pd(L), check_and_make_pd(R)
    return L, R
//  [python]
def sum_kronecker_products(
        left_tensor: Tensor,
        right_tensor: Tensor,
        assert_positive_definite: bool = True) -> Tuple[Tensor, Tensor]:
    r"""Approximation of sums of general Kronecker products with the power method.

    It approximates
        .. math:
            L, R = \text{arg min}_{L, R} \| \frac{1}{B} \sum_{i = 1}^B
                left_tensor[i] \otimes right_tensor[i] - L \otimes R\|

    It corresponds to a generalized version of Algorithm 2 of the appendix of
    `Natural continual learning: success is a journey, not (just) a destination
    <https://proceedings.neurips.cc/paper/2021/file/ec5aa0b7846082a2415f0902f0da88f2-Paper.pdf>_`

    Args:
        left_tensor: (shape [B, M, M]) The first tensor
        right_tensor: (shape [B, N, N]) The second tensor
        assert_positive_definite: If the result should be made positive definite
            to counteract numeric instabilities

    Returns:
        The two tensors L (shape [M, M]) and R (shape [N, N])
    """
    l = left_tensor.view(left_tensor.shape[0], -1)
    r = right_tensor.view(right_tensor.shape[0], -1)
    Q, _ = torch.linalg.qr(l.T)
    H = (l @ Q).T @ r
    U, s, VT = torch.linalg.svd(H, full_matrices=False)
    y = s[0].sqrt() * VT[0]
    x = s[0].sqrt() * Q @ U[:, 0]
    L, R = x.view(*left_tensor.shape[1:]), y.view(*right_tensor.shape[1:])
    if torch.trace(L) < 0 and torch.trace(R) < 0:
        L, R = -L, -R
    if assert_positive_definite:
        L, R = check_and_make_pd(L), check_and_make_pd(R)
    return L, R
//  [python]
def local_encoder_pullback_xt(
        self, x, t, op=None, block_idx=None,
        pca_rank=50, chunk_size=25, min_iter=10, max_iter=100, convergence_threshold=1e-3,
    ):
    '''
    Args
        - x : zt
        - op : ['down', 'mid', 'up']
        - block_idx : op == down, up : [0,1,2,3], op == mid : [0]
        - pooling : ['pixel-sum', 'channel-sum', 'single-channel', 'multiple-channel']
    Returns
        - h : hidden feature
    '''
    num_chunk = pca_rank // chunk_size if pca_rank % chunk_size == 0 else pca_rank // chunk_size + 1

    # get h samples
    time_s = time.time()

    # necessary variables
    h_shape = self.get_h(
        x, t=t, op=op, block_idx=block_idx, 
    ).shape

    c_i, w_i, h_i = x.size(1), x.size(2), x.size(3)
    c_o, w_o, h_o = h_shape[1], h_shape[2], h_shape[3]

    a = torch.tensor(0., device=x.device, dtype=x.dtype)

    # Algorithm 1
    vT = torch.randn(c_i*w_i*h_i, pca_rank, device=x.device, dtype=torch.float)
    q, r = torch.linalg.qr(vT)
    vT = q
    v = vT.T
    v = v.view(-1, c_i, w_i, h_i)

    time_s = time.time()
    for i in range(max_iter):
        v = v.to(device=x.device, dtype=x.dtype)
        v_prev = v.detach().cpu().clone()
        
        
        u = []
        v_buffer = list(v.chunk(num_chunk))
        for vi in v_buffer:
            # g = lambda a : get_h(x + a*vi)

            g = lambda a : self.get_h(
                x + a*vi, t=t, op=op, block_idx=block_idx, 
            )

            ui = torch.func.jacfwd(g, argnums=0, has_aux=False, randomness='error')(a)
            u.append(ui.detach().cpu().clone())
        u = torch.cat(u, dim=0)
        u = u.to(x.device, x.dtype)

        # time_s = time.time()
        # g = lambda a : get_h(x + a*v)
        # u = torch.func.jacfwd(g, argnums=0, has_aux=False, randomness='error')(a)
        # time_e = time.time()
        # print('single vi jacfwd t ==', time_e - time_s)

        g = lambda x : einsum(
            u, self.get_h(
                x, t=t, op=op, block_idx=block_idx, 
            ), 'b c w h, i c w h -> b'
        )
        v_ = torch.autograd.functional.jacobian(g, x)
        v_ = v_.view(-1, c_i*w_i*h_i)

        _, s, v = torch.linalg.svd(v_, full_matrices=False)
        v = v.view(-1, c_i, w_i, h_i)
        u = u.view(-1, c_o, w_o, h_o)
        
        convergence = torch.dist(v_prev, v.detach().cpu()).item()
        print(f'power method : {i}-th step convergence : ', convergence)
        
        if torch.allclose(v_prev, v.detach().cpu(), atol=convergence_threshold) and (i > min_iter):
            print('reach convergence threshold : ', convergence)
            break

    time_e = time.time()
    print('power method runtime ==', time_e - time_s)

    # u, s, vT = u.view(-1, c_o*w_o*h_o).T.detach(), s.sqrt().detach(), v.view(-1, c_i*w_i*h_i).detach()
    u, s, vT = u.reshape(-1, c_o*w_o*h_o).T.detach(), s.sqrt().detach(), v.reshape(-1, c_i*w_i*h_i).detach()
    return u, s, vT