jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def fake_poweriteration(input: torch.Tensor, loop, rank, device, p_base, q_base):
    # input size [batch,num_head,seq_len,model_dim/num_head]
    # -> [batch,seq_len,model_dim] -> [batch * seq_len,model_dim]
    # p_base = torch.rand(input.shape[3] * input.shape[1], rank).to(device)
    # q_base = torch.rand(input.shape[0] * input.shape[2], rank).to(device)
    dtype = input.dtype
    batch, num_head, seq_len, sep_dim = input.shape
    input = (
        input.permute(0, 2, 1, 3).contiguous().view(batch, seq_len, sep_dim * num_head)
    )  # convert to 32bits for qr decomposition
    input = input.view(seq_len * batch, sep_dim * num_head)
    input = input.float()
    if q_base is not None and p_base is not None:
        p_base[0] = p_base[0].float()
        q_base[0] = q_base[0].float()
    else:
        p_base = [torch.rand(sep_dim * num_head, rank).to(input.device)]
        q_base = [torch.rand(batch * seq_len, rank).to(input.device)]
    # 3 calculation = loop * (matmul) + 2 * qrO(n^2)
    for i in range(loop):
        if i == loop - 1:
            q, r = torch.qr(p_base[0])
            p_base[0] = q
        q_base[0] = input @ p_base[0]
        if i == loop - 1:
            q, r = torch.qr(q_base[0])
            q_base[0] = q
        p_base[0] = torch.transpose(input, 0, 1) @ q_base[0]
    input = q_base[0] @ torch.transpose(p_base[0], 0, 1)
    input = input.view(batch, seq_len, num_head, sep_dim)
    input = input.permute(0, 2, 1, 3)
    input = input.type(dtype)
    p_base[0] = p_base[0].half()
    q_base[0] = q_base[0].half()
    return input
//  [python]
def fake_poweriteration(input: torch.Tensor, loop, rank, device, p_base, q_base):
    # input size [batch,num_head,seq_len,model_dim/num_head]
    # -> [batch,seq_len,model_dim] -> [batch * seq_len,model_dim]
    # p_base = torch.rand(input.shape[3] * input.shape[1], rank).to(device)
    # q_base = torch.rand(input.shape[0] * input.shape[2], rank).to(device)
    dtype = input.dtype
    batch, num_head, seq_len, sep_dim = input.shape
    input = (
        input.permute(0, 2, 1, 3).contiguous().view(batch, seq_len, sep_dim * num_head)
    )  # convert to 32bits for qr decomposition
    input = input.view(seq_len * batch, sep_dim * num_head)
    input = input.float()
    if q_base is not None and p_base is not None:
        p_base[0] = p_base[0].float()
        q_base[0] = q_base[0].float()
    else:
        p_base = [torch.rand(sep_dim * num_head, rank).to(input.device)]
        q_base = [torch.rand(batch * seq_len, rank).to(input.device)]
    # 3 calculation = loop * (matmul) + 2 * qrO(n^2)
    for i in range(loop):
        if i == loop - 1:
            p_base[0] = torch.linalg.qr(p_base[0]).Q
        q_base[0] = input @ p_base[0]
        if i == loop - 1:
            q_base[0] = torch.linalg.qr(q_base[0]).Q
        p_base[0] = torch.transpose(input, 0, 1) @ q_base[0]
    input = q_base[0] @ torch.transpose(p_base[0], 0, 1)
    input = input.view(batch, seq_len, num_head, sep_dim)
    input = input.permute(0, 2, 1, 3)
    input = input.type(dtype)
    p_base[0] = p_base[0].half()
    q_base[0] = q_base[0].half()
    return input
//  [python]
def sample_symmetric(dim, mean=0, std=10, **kwargs):
    """generate a random symmetric matrix with a particular eigenvalue
    distribution"""
    # 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 with a particular eigenvalue
    distribution"""
    # 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 factor_orthogonalize(
        self,
        mu: int):

        """
        Pushes the factor's non-orthogonal part to its corresponding core.

        This method works in place.

        :param mu: an int between 0 and N-1
        """

        if self.Us[mu] is None:
            return

        Q, R = torch.qr(self.Us[mu])
        self.Us[mu] = Q

        if self.batch:
            m = 3
            idx1 = 'bjk,baj->bak'
            idx2 = 'bijk,baj->biak'
        else:
            m = 2
            idx1 = 'jk,aj->ak'
            idx2 = 'ijk,aj->iak'

        if self.cores[mu].dim() == m:
            self.cores[mu] = torch.einsum(idx1, (self.cores[mu], R))
        else:
            self.cores[mu] = torch.einsum(idx2, (self.cores[mu], R))
//  [python]
def factor_orthogonalize(
        self,
        mu: int):

        """
        Pushes the factor's non-orthogonal part to its corresponding core.

        This method works in place.

        :param mu: an int between 0 and N-1
        """

        if self.Us[mu] is None:
            return

        Q, R = torch.linalg.qr(self.Us[mu])
        self.Us[mu] = Q

        if self.batch:
            m = 3
            idx1 = 'bjk,baj->bak'
            idx2 = 'bijk,baj->biak'
        else:
            m = 2
            idx1 = 'jk,aj->ak'
            idx2 = 'ijk,aj->iak'

        if self.cores[mu].dim() == m:
            self.cores[mu] = torch.einsum(idx1, (self.cores[mu], R))
        else:
            self.cores[mu] = torch.einsum(idx2, (self.cores[mu], R))
//  [python]
def orthogonal_(tensor, gain=1):
    r"""Fills the input `Tensor` with a (semi) orthogonal matrix, as
    described in `Exact solutions to the nonlinear dynamics of learning in deep
    linear neural networks` - Saxe, A. et al. (2013). The input tensor must have
    at least 2 dimensions, and for tensors with more than 2 dimensions the
    trailing dimensions are flattened.

    Args:
        tensor: an n-dimensional `torch.Tensor`, where :math:`n \geq 2`
        gain: optional scaling factor

    Examples:
        >>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_LAPACK)
        >>> w = torch.empty(3, 5)
        >>> nn.init.orthogonal_(w)
    """
    if tensor.ndimension() < 2:
        raise ValueError("Only tensors with 2 or more dimensions are supported")

    if tensor.numel() == 0:
        # no-op
        return tensor
    rows = tensor.size(0)
    cols = tensor.numel() // rows
    flattened = tensor.new(rows, cols).normal_(0, 1)

    if rows < cols:
        flattened.t_()

    # Compute the qr factorization
    q, r = torch.qr(flattened)
    # Make Q uniform according to https://arxiv.org/pdf/math-ph/0609050.pdf
    d = torch.diag(r, 0)
    ph = d.sign()
    q *= ph

    if rows < cols:
        q.t_()

    with torch.no_grad():
        tensor.view_as(q).copy_(q)
        tensor.mul_(gain)
    return tensor
//  [python]
def orthogonal_(tensor, gain=1):
    r"""Fills the input `Tensor` with a (semi) orthogonal matrix, as
    described in `Exact solutions to the nonlinear dynamics of learning in deep
    linear neural networks` - Saxe, A. et al. (2013). The input tensor must have
    at least 2 dimensions, and for tensors with more than 2 dimensions the
    trailing dimensions are flattened.

    Args:
        tensor: an n-dimensional `torch.Tensor`, where :math:`n \geq 2`
        gain: optional scaling factor

    Examples:
        >>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_LAPACK)
        >>> w = torch.empty(3, 5)
        >>> nn.init.orthogonal_(w)
    """
    if tensor.ndimension() < 2:
        raise ValueError("Only tensors with 2 or more dimensions are supported")

    if tensor.numel() == 0:
        # no-op
        return tensor
    rows = tensor.size(0)
    cols = tensor.numel() // rows
    flattened = tensor.new(rows, cols).normal_(0, 1)

    if rows < cols:
        flattened.t_()

    # Compute the qr factorization
    q, r = torch.linalg.qr(flattened)
    # Make Q uniform according to https://arxiv.org/pdf/math-ph/0609050.pdf
    d = torch.diag(r, 0)
    ph = d.sign()
    q *= ph

    if rows < cols:
        q.t_()

    with torch.no_grad():
        tensor.view_as(q).copy_(q)
        tensor.mul_(gain)
    return tensor
//  [python]
def local_encoder_pullback_xt(
            self, x=None, t=None, op=None, block_idx=None, pca_rank=16, chunk_size=25,
            min_iter=10, max_iter=100, convergence_threshold=1e-3,
        ):
        '''
        Args
            - sample : 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
        '''
        # necessary variables
        num_chunk = pca_rank // chunk_size if pca_rank % chunk_size == 0 else pca_rank // chunk_size + 1
        get_h = lambda x : self.get_h(
            x, t=t, op=op, block_idx=block_idx,
        )
        h_shape = get_h(x).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)

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

        for i in range(max_iter):
            v_prev = v.detach().cpu().clone()

            u = []
            time_s = time.time()

            v_buffer = list(v.chunk(num_chunk))
            for vi in v_buffer:
                # g = lambda a : get_h(x + a*vi.unsqueeze(0) if vi.size(0) == v.size(-1) else x + a*vi)
                g = lambda a : get_h(x + a*vi)
                ui = torch.func.jacfwd(g, argnums=0, has_aux=False, randomness='error')(a)
                u.append(ui.detach().cpu().clone())
            time_e = time.time()
            print('single v jacfwd t ==', time_e - time_s)
            u = torch.cat(u, dim=0)
            u = u.to(x.device)

            # 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, get_h(x), '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())
            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

            if i == max_iter - 1:
                print('last convergence : ', convergence)

        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()
        return u, s, vT
//  [python]
def local_encoder_pullback_xt(
            self, x=None, t=None, op=None, block_idx=None, pca_rank=16, chunk_size=25,
            min_iter=10, max_iter=100, convergence_threshold=1e-3,
        ):
        '''
        Args
            - sample : 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
        '''
        # necessary variables
        num_chunk = pca_rank // chunk_size if pca_rank % chunk_size == 0 else pca_rank // chunk_size + 1
        get_h = lambda x : self.get_h(
            x, t=t, op=op, block_idx=block_idx,
        )
        h_shape = get_h(x).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)

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

        for i in range(max_iter):
            v_prev = v.detach().cpu().clone()

            u = []
            time_s = time.time()

            v_buffer = list(v.chunk(num_chunk))
            for vi in v_buffer:
                # g = lambda a : get_h(x + a*vi.unsqueeze(0) if vi.size(0) == v.size(-1) else x + a*vi)
                g = lambda a : get_h(x + a*vi)
                ui = torch.func.jacfwd(g, argnums=0, has_aux=False, randomness='error')(a)
                u.append(ui.detach().cpu().clone())
            time_e = time.time()
            print('single v jacfwd t ==', time_e - time_s)
            u = torch.cat(u, dim=0)
            u = u.to(x.device)

            # 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, get_h(x), '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())
            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

            if i == max_iter - 1:
                print('last convergence : ', convergence)

        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()
        return u, s, vT
//  [python]
def local_decoder_pullback_xt(
            self, x=None, t=None, op=None, block_idx=None, pca_rank=50, chunk_size=10,
            min_iter=10, max_iter=100, convergence_threshold=1e-3,
        ):
        '''
        Returns
            - u     : h space direction (dim = |dh, pc|)
            - s     : eigenvalue spectrum
            - vT    : e space direction (dim = |pc, de|)
        '''
        # necessary variables
        num_chunk = pca_rank // chunk_size if pca_rank % chunk_size == 0 else pca_rank // chunk_size + 1
        assert op == 'mid'
        assert block_idx == 0
        h = self.get_h(
            x=x, t=t, op=op, block_idx=block_idx,
        )

        get_h_to_e = lambda h : self.get_h_to_e(
            x=x, t=t, input_h=h,
            op=op, block_idx=block_idx, verbose=False,
        )

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

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

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

        for i in range(max_iter):
            # iterate over pca directions to avoid OOM
            v_prev = v.detach().cpu().clone()
            u = []
            time_s = time.time()

            v_buffer = list(v.chunk(num_chunk))
            for vi in v_buffer:
                g = lambda a : get_h_to_e(h + a*vi)
                ui = torch.func.jacfwd(g, argnums=0, has_aux=False, randomness='error')(a)
                u.append(ui.detach().cpu().clone())
            time_e = time.time()
            print('single v jacrev t ==', time_e - time_s)
            u = torch.cat(u, dim=0)
            u = u.to(x.device)
            
            # time_s = time.time()
            # g = lambda a : get_h_to_e(h + a*v)
            # u = torch.func.jacfwd(g, argnums=0, has_aux=False, randomness='error')(a)
            # time_e = time.time()
            # print('single vi jacrev t ==', time_e - time_s)

            g = lambda h : einsum(u, get_h_to_e(h), 'b c w h, i c w h -> b')
            v_ = torch.autograd.functional.jacobian(g, h)
            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())
            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
            else:
                print('convergence ==', torch.dist(v_prev, v.detach().cpu()))
        
        # decoder jac do not return vT
        u, s, vT = v.view(-1, c_i*w_i*h_i).T.detach(), s.sqrt().detach(), u.view(-1, c_o*w_o*h_o).detach()
        return u, s, vT
//  [python]
def local_decoder_pullback_xt(
            self, x=None, t=None, op=None, block_idx=None, pca_rank=50, chunk_size=10,
            min_iter=10, max_iter=100, convergence_threshold=1e-3,
        ):
        '''
        Returns
            - u     : h space direction (dim = |dh, pc|)
            - s     : eigenvalue spectrum
            - vT    : e space direction (dim = |pc, de|)
        '''
        # necessary variables
        num_chunk = pca_rank // chunk_size if pca_rank % chunk_size == 0 else pca_rank // chunk_size + 1
        assert op == 'mid'
        assert block_idx == 0
        h = self.get_h(
            x=x, t=t, op=op, block_idx=block_idx,
        )

        get_h_to_e = lambda h : self.get_h_to_e(
            x=x, t=t, input_h=h,
            op=op, block_idx=block_idx, verbose=False,
        )

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

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

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

        for i in range(max_iter):
            # iterate over pca directions to avoid OOM
            v_prev = v.detach().cpu().clone()
            u = []
            time_s = time.time()

            v_buffer = list(v.chunk(num_chunk))
            for vi in v_buffer:
                g = lambda a : get_h_to_e(h + a*vi)
                ui = torch.func.jacfwd(g, argnums=0, has_aux=False, randomness='error')(a)
                u.append(ui.detach().cpu().clone())
            time_e = time.time()
            print('single v jacrev t ==', time_e - time_s)
            u = torch.cat(u, dim=0)
            u = u.to(x.device)
            
            # time_s = time.time()
            # g = lambda a : get_h_to_e(h + a*v)
            # u = torch.func.jacfwd(g, argnums=0, has_aux=False, randomness='error')(a)
            # time_e = time.time()
            # print('single vi jacrev t ==', time_e - time_s)

            g = lambda h : einsum(u, get_h_to_e(h), 'b c w h, i c w h -> b')
            v_ = torch.autograd.functional.jacobian(g, h)
            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())
            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
            else:
                print('convergence ==', torch.dist(v_prev, v.detach().cpu()))
        
        # decoder jac do not return vT
        u, s, vT = v.view(-1, c_i*w_i*h_i).T.detach(), s.sqrt().detach(), u.view(-1, c_o*w_o*h_o).detach()
        return u, s, vT
//  [python]
def local_x0_decoder_pullback_xt(
            self, x=None, t=None, at=None, op=None, block_idx=None, pca_rank=50, chunk_size=10,
            num_chunk=None, min_iter=10, max_iter=100, convergence_threshold=1e-3,
        ):
        '''
        Returns
            - u     : h space direction (dim = |dh, pc|)
            - s     : eigenvalue spectrum
            - vT    : e space direction (dim = |pc, de|)
        '''
        # necessary variables
        assert op == 'mid'
        assert block_idx == 0
        h = self.get_h(
            x=x, t=t, op=op, block_idx=block_idx,
        )

        get_h_to_x0 = lambda h : (x - (1-at).sqrt() * self.get_h_to_e(
            x=x, t=t, input_h=h,
            op=op, block_idx=block_idx, verbose=False,
        )) / at.sqrt()

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

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

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

        for i in range(max_iter):
            # iterate over pca directions to avoid OOM
            v_prev = v.detach().cpu().clone()
            u = []
            time_s = time.time()

            v_buffer = list(v.chunk(v.size(0) // chunk_size))
            for vi in v_buffer:
                g = lambda a : get_h_to_x0(h + a*vi)
                ui = torch.func.jacfwd(g, argnums=0, has_aux=False, randomness='error')(a)
                u.append(ui.detach().cpu().clone())
            time_e = time.time()
            print('single v jacrev t ==', time_e - time_s)
            u = torch.cat(u, dim=0)
            u = u.to(x.device)
            
            # time_s = time.time()
            # g = lambda a : get_h_to_x0(h + a*v)
            # u = torch.func.jacfwd(g, argnums=0, has_aux=False, randomness='error')(a)
            # time_e = time.time()
            # print('single vi jacrev t ==', time_e - time_s)

            g = lambda h : einsum(u, get_h_to_x0(h), 'b c w h, i c w h -> b')
            v_ = torch.autograd.functional.jacobian(g, h)
            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())
            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
            else:
                print('convergence ==', torch.dist(v_prev, v.detach().cpu()))
        
        # decoder jac do not return vT
        u, s, vT = v.view(-1, c_i*w_i*h_i).T.detach(), s.sqrt().detach(), u.view(-1, c_o*w_o*h_o).detach()
        return u, s, vT
//  [python]
def local_x0_decoder_pullback_xt(
            self, x=None, t=None, at=None, op=None, block_idx=None, pca_rank=50, chunk_size=10,
            num_chunk=None, min_iter=10, max_iter=100, convergence_threshold=1e-3,
        ):
        '''
        Returns
            - u     : h space direction (dim = |dh, pc|)
            - s     : eigenvalue spectrum
            - vT    : e space direction (dim = |pc, de|)
        '''
        # necessary variables
        assert op == 'mid'
        assert block_idx == 0
        h = self.get_h(
            x=x, t=t, op=op, block_idx=block_idx,
        )

        get_h_to_x0 = lambda h : (x - (1-at).sqrt() * self.get_h_to_e(
            x=x, t=t, input_h=h,
            op=op, block_idx=block_idx, verbose=False,
        )) / at.sqrt()

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

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

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

        for i in range(max_iter):
            # iterate over pca directions to avoid OOM
            v_prev = v.detach().cpu().clone()
            u = []
            time_s = time.time()

            v_buffer = list(v.chunk(v.size(0) // chunk_size))
            for vi in v_buffer:
                g = lambda a : get_h_to_x0(h + a*vi)
                ui = torch.func.jacfwd(g, argnums=0, has_aux=False, randomness='error')(a)
                u.append(ui.detach().cpu().clone())
            time_e = time.time()
            print('single v jacrev t ==', time_e - time_s)
            u = torch.cat(u, dim=0)
            u = u.to(x.device)
            
            # time_s = time.time()
            # g = lambda a : get_h_to_x0(h + a*v)
            # u = torch.func.jacfwd(g, argnums=0, has_aux=False, randomness='error')(a)
            # time_e = time.time()
            # print('single vi jacrev t ==', time_e - time_s)

            g = lambda h : einsum(u, get_h_to_x0(h), 'b c w h, i c w h -> b')
            v_ = torch.autograd.functional.jacobian(g, h)
            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())
            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
            else:
                print('convergence ==', torch.dist(v_prev, v.detach().cpu()))
        
        # decoder jac do not return vT
        u, s, vT = v.view(-1, c_i*w_i*h_i).T.detach(), s.sqrt().detach(), u.view(-1, c_o*w_o*h_o).detach()
        return u, s, vT
//  [python]
def left_orthogonalize(
            self,
            mu: int):

        """
        Makes the mu-th core left-orthogonal and pushes the R factor to its right core. This may change the ranks
        of the cores.

        This method works in place.

        Note: internally, this method will turn CP (or CP-Tucker) cores into TT (or TT-Tucker) ones.

        :param mu: an int between 0 and N-1

        :return: the R factor
        """

        assert 0 <= mu < self.dim()-1
        self.factor_orthogonalize(mu)
        Q, R = torch.qr(tn.left_unfolding(self.cores[mu], batch=self.batch))

        if self.batch:
            self.cores[mu] = Q.reshape(self.cores[mu].shape[:-1] + (Q.shape[2], ))
        else:
            self.cores[mu] = Q.reshape(self.cores[mu].shape[:-1] + (Q.shape[1], ))

        rightcoreR = tn.right_unfolding(self.cores[mu + 1], batch=self.batch)

        if self.batch:
            self.cores[mu + 1] = (R @ rightcoreR).reshape((R.shape[0], R.shape[1]) + self.cores[mu + 1].shape[2:])
        else:
            self.cores[mu + 1] = (R @ rightcoreR).reshape((R.shape[0], ) + self.cores[mu + 1].shape[1:])
        return R
//  [python]
def left_orthogonalize(
            self,
            mu: int):

        """
        Makes the mu-th core left-orthogonal and pushes the R factor to its right core. This may change the ranks
        of the cores.

        This method works in place.

        Note: internally, this method will turn CP (or CP-Tucker) cores into TT (or TT-Tucker) ones.

        :param mu: an int between 0 and N-1

        :return: the R factor
        """

        assert 0 <= mu < self.dim()-1
        self.factor_orthogonalize(mu)
        Q, R = torch.linalg.qr(tn.left_unfolding(self.cores[mu], batch=self.batch))

        if self.batch:
            self.cores[mu] = Q.reshape(self.cores[mu].shape[:-1] + (Q.shape[2], ))
        else:
            self.cores[mu] = Q.reshape(self.cores[mu].shape[:-1] + (Q.shape[1], ))

        rightcoreR = tn.right_unfolding(self.cores[mu + 1], batch=self.batch)

        if self.batch:
            self.cores[mu + 1] = (R @ rightcoreR).reshape((R.shape[0], R.shape[1]) + self.cores[mu + 1].shape[2:])
        else:
            self.cores[mu + 1] = (R @ rightcoreR).reshape((R.shape[0], ) + self.cores[mu + 1].shape[1:])
        return R
//  [python]
def synthetic(inputargs):
        # The experiment in section 6.1
        args = {
            "r1":3,
            "r2":3,
            "optim":"SGD",
            "lr":0.1,
            "epochs":500,
            "beta":1e-1,
            "seed":100,
            "verbose":2,
            "wd":0,
            #"epsilon":1e-8,
        }
        print(args)
        d = 120
        ndt = 100
        N = 100
        #Ugt = torch.randn(d,args["r1"])/np.sqrt(d)
        device = "cuda"

        Ugt = torch.zeros(d,args["r1"]).to(device)
        for i in range(args["r1"]):
            Ugt[i,i]=1#0

        Q, R = torch.qr(torch.randn(d,d))    
        Q = Q.to(device)
        Ult = torch.zeros(N, d, args["r2"]).to(device)
        ter_id = 119
        ini_id = 5
        for clid in range(N):
            for cpid in range(args["r2"]):
                Ult[clid, ((cpid + clid) % (ter_id - ini_id + 1) + ini_id), cpid] = 1
        with torch.no_grad():
            Ugt = Q@Ugt
            for i in range(N):
                Ult[i,:,:] = Q@Ult[i,:,:]
        import torch.nn.functional as F
        pavg = sum(Ult[i,:,:]@Ult[i,:,:].T for i in range(N))/N
        up,sp,vhp = torch.svd(pavg)
        #print(sp)
        print("theta is %.6f"%(1-sp[0]))
        
        
        Y = []
        ulst = []
        vgst = []
        vlst = []
        for i in range(N):
            Vgt = torch.randn(ndt,args["r1"]).to(device)#*100
            Qvg,R = torch.qr(Vgt)
            Vgt = Qvg
            Ult[i,:,:] -= Ugt @ torch.inverse(Ugt.T@Ugt) @ Ugt.T @ Ult[i,:,:]
            ulst.append(Ult[i,:,:])
            Vlt = torch.randn(ndt,args["r2"]).to(device)#*100
            Vlt -= Vgt @ torch.inverse(Vgt.T@Vgt) @ Vgt.T@Vlt
            Qvl,R = torch.qr(Vlt)
            Vlt = Qvl
            
            ygpart = Ugt@Vgt.T           

            yipart = Ult[i,:,:]@Vlt.T           
            
            Y.append(ygpart+yipart)
            vgst.append(Vgt)
            vlst.append(Vlt)
        

        args["global_subspace_err_metric"]= lambda x: subspace_error(Ugt,x)
        args["local_subspace_err_metric"] = lambda x: sum([subspace_error(ulst[i],x[i]) for i in range(N)])/N
     
        args["global_recovery_error"] = lambda x: g_recovery([Ugt,vgst,ulst,vlst],x,'f')
        args["local_recovery_error"] = lambda x: l_recovery([Ugt,vgst,ulst,vlst],x,'f')

        Ug,Vg,Ul,Vl = heterogeneous_matrix_factorization(Y,args,verbose=20)
//  [python]
def synthetic(inputargs):
        # The experiment in section 6.1
        args = {
            "r1":3,
            "r2":3,
            "optim":"SGD",
            "lr":0.1,
            "epochs":500,
            "beta":1e-1,
            "seed":100,
            "verbose":2,
            "wd":0,
            #"epsilon":1e-8,
        }
        print(args)
        d = 120
        ndt = 100
        N = 100
        #Ugt = torch.randn(d,args["r1"])/np.sqrt(d)
        device = "cuda"

        Ugt = torch.zeros(d,args["r1"]).to(device)
        for i in range(args["r1"]):
            Ugt[i,i]=1#0

        Q, R = torch.linalg.qr(torch.randn(d,d))    
        Q = Q.to(device)
        Ult = torch.zeros(N, d, args["r2"]).to(device)
        ter_id = 119
        ini_id = 5
        for clid in range(N):
            for cpid in range(args["r2"]):
                Ult[clid, ((cpid + clid) % (ter_id - ini_id + 1) + ini_id), cpid] = 1
        with torch.no_grad():
            Ugt = Q@Ugt
            for i in range(N):
                Ult[i,:,:] = Q@Ult[i,:,:]
        import torch.nn.functional as F
        pavg = sum(Ult[i,:,:]@Ult[i,:,:].T for i in range(N))/N
        up,sp,vhp = torch.linalg.svd(pavg)
        #print(sp)
        print("theta is %.6f"%(1-sp[0]))
        
        
        Y = []
        ulst = []
        vgst = []
        vlst = []
        for i in range(N):
            Vgt = torch.randn(ndt,args["r1"]).to(device)#*100
            Qvg,R = torch.linalg.qr(Vgt)
            Vgt = Qvg
            Ult[i,:,:] -= Ugt @ torch.inverse(Ugt.T@Ugt) @ Ugt.T @ Ult[i,:,:]
            ulst.append(Ult[i,:,:])
            Vlt = torch.randn(ndt,args["r2"]).to(device)#*100
            Vlt -= Vgt @ torch.inverse(Vgt.T@Vgt) @ Vgt.T@Vlt
            Qvl,R = torch.linalg.qr(Vlt)
            Vlt = Qvl
            
            ygpart = Ugt@Vgt.T           

            yipart = Ult[i,:,:]@Vlt.T           
            
            Y.append(ygpart+yipart)
            vgst.append(Vgt)
            vlst.append(Vlt)
        

        args["global_subspace_err_metric"]= lambda x: subspace_error(Ugt,x)
        args["local_subspace_err_metric"] = lambda x: sum([subspace_error(ulst[i],x[i]) for i in range(N)])/N
     
        args["global_recovery_error"] = lambda x: g_recovery([Ugt,vgst,ulst,vlst],x,'f')
        args["local_recovery_error"] = lambda x: l_recovery([Ugt,vgst,ulst,vlst],x,'f')

        Ug,Vg,Ul,Vl = heterogeneous_matrix_factorization(Y,args,verbose=20)
//  [python]
def right_orthogonalize(
            self,
            mu: int):

        """
        Makes the mu-th core right-orthogonal and pushes the L factor to its left core. Note: this may change the ranks
         of the tensor.

        This method works in place.

        Note: internally, this method will turn CP (or CP-Tucker) cores into TT (or TT-Tucker) ones.

        :param mu: an int between 0 and N-1

        :return: the L factor
        """

        assert 1 <= mu < self.dim()
        self.factor_orthogonalize(mu)
        # Torch has no rq() decomposition
        if self.batch:
            Q, L = torch.qr(tn.right_unfolding(self.cores[mu], batch=self.batch).permute(0, 2, 1))
            L = L.permute(0, 2, 1)
            Q = Q.permute(0, 2, 1)
        else:
            Q, L = torch.qr(tn.right_unfolding(self.cores[mu], batch=self.batch).permute(1, 0))
            L = L.permute(1, 0)
            Q = Q.permute(1, 0)

        if self.batch:
            self.cores[mu] = Q.reshape((Q.shape[:2]) + self.cores[mu].shape[2:])
        else:
            self.cores[mu] = Q.reshape((Q.shape[0], ) + self.cores[mu].shape[1:])

        leftcoreL = tn.left_unfolding(self.cores[mu-1], batch=self.batch)
        if self.batch:
            self.cores[mu - 1] = (leftcoreL @ L).reshape(self.cores[mu-1].shape[:-1] + (L.shape[2], ))
        else:
            self.cores[mu - 1] = (leftcoreL @ L).reshape(self.cores[mu-1].shape[:-1] + (L.shape[1], ))
        return L
//  [python]
def right_orthogonalize(
            self,
            mu: int):

        """
        Makes the mu-th core right-orthogonal and pushes the L factor to its left core. Note: this may change the ranks
         of the tensor.

        This method works in place.

        Note: internally, this method will turn CP (or CP-Tucker) cores into TT (or TT-Tucker) ones.

        :param mu: an int between 0 and N-1

        :return: the L factor
        """

        assert 1 <= mu < self.dim()
        self.factor_orthogonalize(mu)
        # Torch has no rq() decomposition
        if self.batch:
            Q, L = torch.linalg.qr(tn.right_unfolding(self.cores[mu], batch=self.batch).permute(0, 2, 1))
            L = L.permute(0, 2, 1)
            Q = Q.permute(0, 2, 1)
        else:
            Q, L = torch.linalg.qr(tn.right_unfolding(self.cores[mu], batch=self.batch).permute(1, 0))
            L = L.permute(1, 0)
            Q = Q.permute(1, 0)

        if self.batch:
            self.cores[mu] = Q.reshape((Q.shape[:2]) + self.cores[mu].shape[2:])
        else:
            self.cores[mu] = Q.reshape((Q.shape[0], ) + self.cores[mu].shape[1:])

        leftcoreL = tn.left_unfolding(self.cores[mu-1], batch=self.batch)
        if self.batch:
            self.cores[mu - 1] = (leftcoreL @ L).reshape(self.cores[mu-1].shape[:-1] + (L.shape[2], ))
        else:
            self.cores[mu - 1] = (leftcoreL @ L).reshape(self.cores[mu-1].shape[:-1] + (L.shape[1], ))
        return L
//  [python]
def uniform_init_(tensor):
    r"""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
    """
    # We re-implement torch.nn.init.orthogonal_, as their treatment of batches
    # is not in a per-matrix base
    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)

        # Make uniform (diag r >= 0)
        d = r.diagonal(dim1=-2, dim2=-1).sign()
        q *= d.unsqueeze(-2)
        if transpose:
            q.transpose_(-2, -1)

        # Make them have positive determinant by multiplying the
        # first column by -1 (does not change the measure)
        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):
    r"""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
    """
    # We re-implement torch.nn.init.orthogonal_, as their treatment of batches
    # is not in a per-matrix base
    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)

        # Make uniform (diag r >= 0)
        d = r.diagonal(dim1=-2, dim2=-1).sign()
        q *= d.unsqueeze(-2)
        if transpose:
            q.transpose_(-2, -1)

        # Make them have positive determinant by multiplying the
        # first column by -1 (does not change the measure)
        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 _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 sample_v(x, hutchinson_samples, orthogonalize=True):
    if hutchinson_samples > x.shape[-1]:
        raise ValueError(
            f"Too many Hutchinson samples: got {hutchinson_samples}, expected <= {x.shape[-1]}"
        )
    v = torch.randn(*x.shape, hutchinson_samples, device=x.device, dtype=x.dtype)
    q = (
        torch.qr(v).Q
        if orthogonalize
        else v / torch.linalg.norm(v, axis=1).unsqueeze(1)
    )
    return q * sqrt(x.shape[-1])
//  [python]
def sample_v(x, hutchinson_samples, orthogonalize=True):
    if hutchinson_samples > x.shape[-1]:
        raise ValueError(
            f"Too many Hutchinson samples: got {hutchinson_samples}, expected <= {x.shape[-1]}"
        )
    v = torch.randn(*x.shape, hutchinson_samples, device=x.device, dtype=x.dtype)
    q = (
        torch.linalg.qr(v).Q
        if orthogonalize
        else v / torch.linalg.norm(v, axis=1).unsqueeze(1)
    )
    return q * sqrt(x.shape[-1])
//  [python]
def orthogonal_(gen, tensor, gain=1):
    r"""Fills the input `Tensor` with a (semi) orthogonal matrix, as
    described in `Exact solutions to the nonlinear dynamics of learning in deep
    linear neural networks` - Saxe, A. et al. (2013). The input tensor must have
    at least 2 dimensions, and for tensors with more than 2 dimensions the
    trailing dimensions are flattened.

    Args:
        tensor: an n-dimensional `torch.Tensor`, where :math:`n \geq 2`
        gain: optional scaling factor

    Examples:
        >>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_LAPACK)
        >>> w = torch.empty(3, 5)
        >>> nn.init.orthogonal_(w)
    """
    if tensor.ndimension() < 2:
        raise ValueError("Only tensors with 2 or more dimensions are supported")

    if tensor.numel() == 0:
        # no-op
        return tensor
    rows = tensor.size(0)
    cols = tensor.numel() // rows
    flattened = tensor.new(rows, cols).normal_(0, 1, generator=gen)

    if rows < cols:
        flattened.t_()

    # Compute the qr factorization
    q, r = torch.qr(flattened)
    # Make Q uniform according to https://arxiv.org/pdf/math-ph/0609050.pdf
    d = torch.diag(r, 0)
    ph = d.sign()
    q *= ph

    if rows < cols:
        q.t_()

    with torch.no_grad():
        tensor.view_as(q).copy_(q)
        tensor.mul_(gain)
    return tensor
//  [python]
def orthogonal_(gen, tensor, gain=1):
    r"""Fills the input `Tensor` with a (semi) orthogonal matrix, as
    described in `Exact solutions to the nonlinear dynamics of learning in deep
    linear neural networks` - Saxe, A. et al. (2013). The input tensor must have
    at least 2 dimensions, and for tensors with more than 2 dimensions the
    trailing dimensions are flattened.

    Args:
        tensor: an n-dimensional `torch.Tensor`, where :math:`n \geq 2`
        gain: optional scaling factor

    Examples:
        >>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_LAPACK)
        >>> w = torch.empty(3, 5)
        >>> nn.init.orthogonal_(w)
    """
    if tensor.ndimension() < 2:
        raise ValueError("Only tensors with 2 or more dimensions are supported")

    if tensor.numel() == 0:
        # no-op
        return tensor
    rows = tensor.size(0)
    cols = tensor.numel() // rows
    flattened = tensor.new(rows, cols).normal_(0, 1, generator=gen)

    if rows < cols:
        flattened.t_()

    # Compute the qr factorization
    q, r = torch.linalg.qr(flattened)
    # Make Q uniform according to https://arxiv.org/pdf/math-ph/0609050.pdf
    d = torch.diag(r, 0)
    ph = d.sign()
    q *= ph

    if rows < cols:
        q.t_()

    with torch.no_grad():
        tensor.view_as(q).copy_(q)
        tensor.mul_(gain)
    return tensor
//  [python]
def _orth_by_qr(
    matrix: torch.Tensor, rows_to_orthogonalize: torch.Tensor
) -> torch.Tensor:
    """Orthogonalize a wavelet matrix through qr decomposition.

    A dense qr-decomposition is used for GPU-efficiency reasons.
    If memory becomes a constraint choose _orth_by_gram_schmidt
    instead, which is implemented using only sparse function calls.

    Args:
        matrix (torch.Tensor): The matrix to orthogonalize.
        rows_to_orthogonalize (torch.Tensor): The matrix rows, which need work.

    Returns:
        torch.Tensor: The corrected sparse matrix.
    """
    selection_indices = torch.stack(
        [
            torch.arange(len(rows_to_orthogonalize), device=matrix.device),
            rows_to_orthogonalize,
        ],
        0,
    )
    selection_matrix = torch.sparse_coo_tensor(
        selection_indices,
        torch.ones_like(rows_to_orthogonalize),
        dtype=matrix.dtype,
        device=matrix.device,
    )

    sel = torch.sparse.mm(selection_matrix, matrix)
    q, _ = torch.qr(sel.to_dense().T)
    q_rows = q.T.to_sparse()

    diag_indices = torch.arange(matrix.shape[0])
    diag = torch.ones_like(diag_indices)
    diag[rows_to_orthogonalize] = 0
    removal_matrix = torch.sparse_coo_tensor(
        torch.stack([diag_indices] * 2, 0),
        diag,
        size=matrix.shape,
        device=matrix.device,
        dtype=matrix.dtype,
    )
    result = torch.sparse.mm(removal_matrix, matrix)
    for pos, row in enumerate(q_rows):
        row = row.unsqueeze(0).coalesce()
        addition_matrix = torch.sparse_coo_tensor(
            torch.stack(
                (row.indices()[0, :] + rows_to_orthogonalize[pos], row.indices()[1, :]),
                0,
            ),
            row.values(),
            size=matrix.shape,
            device=matrix.device,
            dtype=matrix.dtype,
        )
        result += addition_matrix
    return result.coalesce()
//  [python]
def _orth_by_qr(
    matrix: torch.Tensor, rows_to_orthogonalize: torch.Tensor
) -> torch.Tensor:
    """Orthogonalize a wavelet matrix through qr decomposition.

    A dense qr-decomposition is used for GPU-efficiency reasons.
    If memory becomes a constraint choose _orth_by_gram_schmidt
    instead, which is implemented using only sparse function calls.

    Args:
        matrix (torch.Tensor): The matrix to orthogonalize.
        rows_to_orthogonalize (torch.Tensor): The matrix rows, which need work.

    Returns:
        torch.Tensor: The corrected sparse matrix.
    """
    selection_indices = torch.stack(
        [
            torch.arange(len(rows_to_orthogonalize), device=matrix.device),
            rows_to_orthogonalize,
        ],
        0,
    )
    selection_matrix = torch.sparse_coo_tensor(
        selection_indices,
        torch.ones_like(rows_to_orthogonalize),
        dtype=matrix.dtype,
        device=matrix.device,
    )

    sel = torch.sparse.mm(selection_matrix, matrix)
    q, _ = torch.linalg.qr(sel.to_dense().T)
    q_rows = q.T.to_sparse()

    diag_indices = torch.arange(matrix.shape[0])
    diag = torch.ones_like(diag_indices)
    diag[rows_to_orthogonalize] = 0
    removal_matrix = torch.sparse_coo_tensor(
        torch.stack([diag_indices] * 2, 0),
        diag,
        size=matrix.shape,
        device=matrix.device,
        dtype=matrix.dtype,
    )
    result = torch.sparse.mm(removal_matrix, matrix)
    for pos, row in enumerate(q_rows):
        row = row.unsqueeze(0).coalesce()
        addition_matrix = torch.sparse_coo_tensor(
            torch.stack(
                (row.indices()[0, :] + rows_to_orthogonalize[pos], row.indices()[1, :]),
                0,
            ),
            row.values(),
            size=matrix.shape,
            device=matrix.device,
            dtype=matrix.dtype,
        )
        result += addition_matrix
    return result.coalesce()