jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def Qdr(a):
    # v0.3 OK
    """
    The function performs the qdr decomposition

        :math:`a = q \cdot d \cdot r`

    to input UniTensor. The UniTensor should be rank-2, untagged. with eachbond's dim should be >=2.

    Args:
        a :
            UniTensor [untagged], rank-2, 1 inbond 1 outbond.

    Return: q , r
        q :
            UniTensor [untagged], rank-2, 1 inbond 1 outbond, the unitary matrix

        d :
            The diagonal matrix [untagged]. It is a diagonal rank-2 UniTensor with 1 inbond 1 outbond and is_diag=True.
        r :
            UniTensor [untagged], rank-2, 1 inbond 1 outbond, the upper triangular matrix
    """
    if isinstance(a,UniTensor):

        ## Qnum_ipoint
        if a.is_symm:
            raise Exception("Qdr(a)","[Abort] curretly don't support symm tensor.")

        if a.is_diag:
            raise Exception("Qr(UniTensor)","[Aboart] Currently not support diagonal tensors.")

        if a.braket is not None:
            raise Exception("Qr(UniTensor)","can only operate on [regular](untagged) tensor")
            
        if a.rowrank != 1:
            raise Exception("Qdr(a)","Should have 1 inbond 1 outbond")

        
        q, r = torch.qr(a.Storage)
        d = r.diag()
        r = (r.t()/d).t()

        tmp = np.argwhere(a.labels<0)
        if len(tmp) == 0:
            tmp = 0
        else:
            tmp = np.min(tmp)

        tq = UniTensor(bonds =[Bond(q.shape[0]),Bond(q.shape[1])],\
                      rowrank=1,\
                      labels=[a.labels[0],tmp-1],\
                      check=False)
        tq._mac(torch_tensor = q)

        td = UniTensor(bonds =[Bond(d.shape[0]),Bond(d.shape[0])],\
                      rowrank=1,\
                      labels=[tmp-1,tmp-2],\
                      is_diag=True,
                      check=False)
        td._mac(torch_tensor = d)

        tr = UniTensor(bonds =[Bond(r.shape[0]),Bond(r.shape[1])],\
                      rowrank=1,\
                      labels=[td.labels[1],a.labels[1]],\
                      check=False)
        tr._mac(torch_tensor = r)

        return tq,td,tr
    else:
        raise Exception("Qdr(UniTensor)","[ERROR] Qdr can only accept UniTensor")
//  [python]
def Qdr(a):
    # v0.3 OK
    """
    The function performs the qdr decomposition

        :math:`a = q \cdot d \cdot r`

    to input UniTensor. The UniTensor should be rank-2, untagged. with eachbond's dim should be >=2.

    Args:
        a :
            UniTensor [untagged], rank-2, 1 inbond 1 outbond.

    Return: q , r
        q :
            UniTensor [untagged], rank-2, 1 inbond 1 outbond, the unitary matrix

        d :
            The diagonal matrix [untagged]. It is a diagonal rank-2 UniTensor with 1 inbond 1 outbond and is_diag=True.
        r :
            UniTensor [untagged], rank-2, 1 inbond 1 outbond, the upper triangular matrix
    """
    if isinstance(a,UniTensor):

        ## Qnum_ipoint
        if a.is_symm:
            raise Exception("Qdr(a)","[Abort] curretly don't support symm tensor.")

        if a.is_diag:
            raise Exception("Qr(UniTensor)","[Aboart] Currently not support diagonal tensors.")

        if a.braket is not None:
            raise Exception("Qr(UniTensor)","can only operate on [regular](untagged) tensor")
            
        if a.rowrank != 1:
            raise Exception("Qdr(a)","Should have 1 inbond 1 outbond")

        
        q, r = torch.linalg.qr(a.Storage)
        d = r.diag()
        r = (r.t()/d).t()

        tmp = np.argwhere(a.labels<0)
        if len(tmp) == 0:
            tmp = 0
        else:
            tmp = np.min(tmp)

        tq = UniTensor(bonds =[Bond(q.shape[0]),Bond(q.shape[1])],\
                      rowrank=1,\
                      labels=[a.labels[0],tmp-1],\
                      check=False)
        tq._mac(torch_tensor = q)

        td = UniTensor(bonds =[Bond(d.shape[0]),Bond(d.shape[0])],\
                      rowrank=1,\
                      labels=[tmp-1,tmp-2],\
                      is_diag=True,
                      check=False)
        td._mac(torch_tensor = d)

        tr = UniTensor(bonds =[Bond(r.shape[0]),Bond(r.shape[1])],\
                      rowrank=1,\
                      labels=[td.labels[1],a.labels[1]],\
                      check=False)
        tr._mac(torch_tensor = r)

        return tq,td,tr
    else:
        raise Exception("Qdr(UniTensor)","[ERROR] Qdr can only accept UniTensor")
//  [python]
def __init__(self, in_channel):
        super().__init__()

        weight = torch.randn(in_channel, in_channel)
        q, _ = torch.qr(weight)
        weight = q.unsqueeze(2).unsqueeze(3)
        self.weight = nn.Parameter(weight)
//  [python]
def __init__(self, in_channel):
        super().__init__()

        weight = torch.randn(in_channel, in_channel)
        q, _ = torch.linalg.qr(weight)
        weight = q.unsqueeze(2).unsqueeze(3)
        self.weight = nn.Parameter(weight)
//  [python]
def projection_matrix_to_CRT_kitti(proj):
    """Split projection matrix of kitti.

    P = C @ [R|T]
    C is upper triangular matrix, so we need to inverse CR and use QR
    stable for all kitti camera projection matrix.

    Args:
        proj (p.array, shape=[4, 4]): Intrinsics of camera.

    Returns:
        tuple[torch.ndarray]: Splited matrix of C, R and T.
    """

    CR = proj[0:3, 0:3]
    CT = proj[0:3, 3]
    RinvCinv = torch.linalg.inv(CR)
    Rinv, Cinv = torch.qr(RinvCinv)
    C = torch.linalg.inv(Cinv)
    R = torch.linalg.inv(Rinv)
    T = Cinv @ CT
    return C, R, T
//  [python]
def projection_matrix_to_CRT_kitti(proj):
    """Split projection matrix of kitti.

    P = C @ [R|T]
    C is upper triangular matrix, so we need to inverse CR and use QR
    stable for all kitti camera projection matrix.

    Args:
        proj (p.array, shape=[4, 4]): Intrinsics of camera.

    Returns:
        tuple[torch.ndarray]: Splited matrix of C, R and T.
    """

    CR = proj[0:3, 0:3]
    CT = proj[0:3, 3]
    RinvCinv = torch.linalg.inv(CR)
    Rinv, Cinv = torch.linalg.qr(RinvCinv)
    C = torch.linalg.inv(Cinv)
    R = torch.linalg.inv(Rinv)
    T = Cinv @ CT
    return C, R, T
//  [python]
def __init__(self, c):
        super().__init__()
        self.conv = torch.nn.Conv1d(c, c, kernel_size=1, stride=1, padding=0, bias=False)

        # Sample a random orthonormal matrix to initialize weights
        W = torch.qr(torch.FloatTensor(c, c).normal_())[0]

        # Ensure determinant is 1.0 not -1.0
        if torch.det(W) < 0:
            W[:, 0] = -1 * W[:, 0]
        W = W.view(c, c, 1)
        self.conv.weight.data = W
        self.inv_conv = None
//  [python]
def __init__(self, c):
        super().__init__()
        self.conv = torch.nn.Conv1d(c, c, kernel_size=1, stride=1, padding=0, bias=False)

        # Sample a random orthonormal matrix to initialize weights
        W = torch.linalg.qr(torch.FloatTensor(c, c).normal_())[0]

        # Ensure determinant is 1.0 not -1.0
        if torch.det(W) < 0:
            W[:, 0] = -1 * W[:, 0]
        W = W.view(c, c, 1)
        self.conv.weight.data = W
        self.inv_conv = None
//  [python]
def forward(self, graph):
        x, edge_index, batch = graph.x, graph.edge_index, graph.batch
        unpool_info = []
        x_info=[]
        cluster_info=[]
        edge_info=[]
        while x.size()[0] > self.coarsening_threshold:
            cluster = graclus(edge_index,num_nodes=x.shape[0])
            cluster_info.append(cluster)
            edge_info.append(edge_index)
            gc = avg_pool(cluster, Batch(batch=batch, x=x, edge_index=edge_index))
            x, edge_index, batch = gc.x, gc.edge_index, gc.batch
        # coarse iterations
        x=torch.eye(2).to(device)
        x=self.conv_coarse(x,edge_index)
        x=self.activation(x)
        while edge_info:
            # un-pooling / interpolation / prolongation / refinement
            edge_index = edge_info.pop()
            output, inverse = torch.unique(cluster_info.pop(), return_inverse=True)
            x = x[inverse]
            # post-smoothing
            for i in range(self.post):
                x = self.activation(self.conv_post[i](x, edge_index))
        x=self.lins1(x)
        x=self.activation(x)
        x=self.lins2(x)
        x=self.activation(x)
        x=self.lins3(x)
        x=self.activation(x)
        x=self.final(x)
        x,_=torch.qr(x,some=True)
        return x
//  [python]
def forward(self, graph):
        x, edge_index, batch = graph.x, graph.edge_index, graph.batch
        unpool_info = []
        x_info=[]
        cluster_info=[]
        edge_info=[]
        while x.size()[0] > self.coarsening_threshold:
            cluster = graclus(edge_index,num_nodes=x.shape[0])
            cluster_info.append(cluster)
            edge_info.append(edge_index)
            gc = avg_pool(cluster, Batch(batch=batch, x=x, edge_index=edge_index))
            x, edge_index, batch = gc.x, gc.edge_index, gc.batch
        # coarse iterations
        x=torch.eye(2).to(device)
        x=self.conv_coarse(x,edge_index)
        x=self.activation(x)
        while edge_info:
            # un-pooling / interpolation / prolongation / refinement
            edge_index = edge_info.pop()
            output, inverse = torch.unique(cluster_info.pop(), return_inverse=True)
            x = x[inverse]
            # post-smoothing
            for i in range(self.post):
                x = self.activation(self.conv_post[i](x, edge_index))
        x=self.lins1(x)
        x=self.activation(x)
        x=self.lins2(x)
        x=self.activation(x)
        x=self.lins3(x)
        x=self.activation(x)
        x=self.final(x)
        x,_=torch.linalg.qr(x,mode='reduced')
        return x
//  [python]
def __init__(
        self,
        size,
        dtype,
    ):
        super().__init__()
        self.size = size

        matrix = torch.randn((size, size), dtype=dtype)
        q, _ = torch.qr(matrix)
        self.weight = nn.Parameter(q)
//  [python]
def __init__(
        self,
        size,
        dtype,
    ):
        super().__init__()
        self.size = size

        matrix = torch.randn((size, size), dtype=dtype)
        q, _ = torch.linalg.qr(matrix)
        self.weight = nn.Parameter(q)
//  [python]
def _rand_haar_slow(n_qubits: int) -> Tensor:
    """
    Detailed in https://arxiv.org/pdf/math-ph/0609050.pdf.

    Textbook implementation, but very expensive. For 12 qubits it takes several seconds.
    For 1 qubit it seems to produce the same distribution as the measure above.
    """
    N = 2**n_qubits
    A = torch.zeros(N, N, dtype=DTYPE).normal_(0, 1)
    B = torch.zeros(N, N, dtype=DTYPE).normal_(0, 1)
    Z = A + 1.0j * B
    Q, R = torch.qr(Z)
    Lambda = torch.diag(torch.diag(R) / torch.diag(R).abs())
    haar_unitary = torch.matmul(Q, Lambda)
    return torch.matmul(haar_unitary, zero_state(n_qubits).squeeze(0)).unsqueeze(0)
//  [python]
def _rand_haar_slow(n_qubits: int) -> Tensor:
    """
    Detailed in https://arxiv.org/pdf/math-ph/0609050.pdf.

    Textbook implementation, but very expensive. For 12 qubits it takes several seconds.
    For 1 qubit it seems to produce the same distribution as the measure above.
    """
    N = 2**n_qubits
    A = torch.zeros(N, N, dtype=DTYPE).normal_(0, 1)
    B = torch.zeros(N, N, dtype=DTYPE).normal_(0, 1)
    Z = A + 1.0j * B
    Q, R = torch.linalg.qr(Z)
    Lambda = torch.diag(torch.diag(R) / torch.diag(R).abs())
    haar_unitary = torch.matmul(Q, Lambda)
    return torch.matmul(haar_unitary, zero_state(n_qubits).squeeze(0)).unsqueeze(0)
//  [python]
def maxKSWDdirection(x, x2='gaussian', weight=None, K=None, maxiter=200,  p=2, eps=1e-6, ATi=None):

    #if x2 is None, find the direction of max sliced Wasserstein distance between x and gaussian
    #if x2 is not None, it needs to have the same shape as x

    if x2 != 'gaussian':
        assert x.shape[1] == x2.shape[1]
        assert weight is None
        if x2.shape[0] > x.shape[0]:
            x2 = x2[torch.randperm(x2.shape[0])][:x.shape[0]]
        elif x2.shape[0] < x.shape[0]:
            x = x[torch.randperm(x.shape[0])][:x2.shape[0]]
    elif weight is not None:
        assert len(weight) == len(x)
        pg = None
        weight = weight / torch.sum(weight)
    else:
        pg = Gaussian_ppf(len(x), device=x.device)

    ndim = x.shape[1]
    if K is None:
        K = ndim

    
    #initialize A. algorithm from https://arxiv.org/pdf/math-ph/0609050.pdf
    if ATi is None:
        ATi = torch.randn(ndim, K, device=x.device)
    else:
        assert ATi.shape[0] == ndim and ATi.shape[1] == K 
    Q, R = torch.qr(ATi)
    L = torch.sign(torch.diag(R))
    AT = (Q * L).T

    lr = 0.1
    down_fac = 0.5
    up_fac = 1.5
    c = 0.5
    
    #algorithm from http://noodle.med.yale.edu/~hdtag/notes/steifel_notes.pdf
    #note that here A = X
    #use backtracking line search
    AT1 = AT.clone()
    for i in range(maxiter):
        AT.requires_grad_(True)
        if x2 == 'gaussian':
            loss = -ObjectiveG(AT @ x.T, pg, p, w=weight)
        else:
            loss = -Objective(AT @ x.T, AT @ x2.T, p)
        loss1 = loss
        GT = torch.autograd.grad(loss, AT)[0]
        AT.requires_grad_(False)
        with torch.no_grad():
            BT = AT.T @ GT - GT.T @ AT
            e = - AT @ BT #dw/dlr
            m = torch.sum(GT * e) #dloss/dlr

            lr /= down_fac
            while loss1 > loss + c*m*lr:
                lr *= down_fac
                if 2*K < ndim:
                    UT = torch.cat((GT, AT), dim=0).double()
                    V = torch.cat((AT.T, -GT.T), dim=1).double()
                    AT1 = (AT.double() - lr * AT.double() @ V @ torch.pinverse(torch.eye(2*K, dtype=torch.double, device=x.device)+lr/2*UT@V) @ UT).to(torch.get_default_dtype())
                else:
                    AT1 = (AT.double() @ (torch.eye(ndim, dtype=torch.double, device=x.device)-lr/2*BT.double()) @ torch.pinverse(torch.eye(ndim, dtype=torch.double, device=x.device)+lr/2*BT.double())).to(torch.get_default_dtype())
            
                if x2 == 'gaussian':
                    loss1 = -ObjectiveG(AT1 @ x.T, pg, p, w=weight)
                else:
                    loss1 = -Objective(AT1 @ x.T, AT1 @ x2.T, p)
        
            if torch.max(torch.abs(AT1-AT)) < eps:
                AT = AT1
                break
        
            lr *= up_fac
            AT = AT1

    if x2 == 'gaussian':
        WD = ObjectiveG(AT @ x.T, pg, p, w=weight, perdim=False)
    else:
        WD = Objective(AT @ x.T, AT @ x2.T, p, perdim=False)
    return AT.T, WD**(1/p)
//  [python]
def maxKSWDdirection(x, x2='gaussian', weight=None, K=None, maxiter=200,  p=2, eps=1e-6, ATi=None):

    #if x2 is None, find the direction of max sliced Wasserstein distance between x and gaussian
    #if x2 is not None, it needs to have the same shape as x

    if x2 != 'gaussian':
        assert x.shape[1] == x2.shape[1]
        assert weight is None
        if x2.shape[0] > x.shape[0]:
            x2 = x2[torch.randperm(x2.shape[0])][:x.shape[0]]
        elif x2.shape[0] < x.shape[0]:
            x = x[torch.randperm(x.shape[0])][:x2.shape[0]]
    elif weight is not None:
        assert len(weight) == len(x)
        pg = None
        weight = weight / torch.sum(weight)
    else:
        pg = Gaussian_ppf(len(x), device=x.device)

    ndim = x.shape[1]
    if K is None:
        K = ndim

    
    #initialize A. algorithm from https://arxiv.org/pdf/math-ph/0609050.pdf
    if ATi is None:
        ATi = torch.randn(ndim, K, device=x.device)
    else:
        assert ATi.shape[0] == ndim and ATi.shape[1] == K 
    Q, R = torch.linalg.qr(ATi)
    L = torch.sign(torch.diag(R))
    AT = (Q * L).T

    lr = 0.1
    down_fac = 0.5
    up_fac = 1.5
    c = 0.5
    
    #algorithm from http://noodle.med.yale.edu/~hdtag/notes/steifel_notes.pdf
    #note that here A = X
    #use backtracking line search
    AT1 = AT.clone()
    for i in range(maxiter):
        AT.requires_grad_(True)
        if x2 == 'gaussian':
            loss = -ObjectiveG(AT @ x.T, pg, p, w=weight)
        else:
            loss = -Objective(AT @ x.T, AT @ x2.T, p)
        loss1 = loss
        GT = torch.autograd.grad(loss, AT)[0]
        AT.requires_grad_(False)
        with torch.no_grad():
            BT = AT.T @ GT - GT.T @ AT
            e = - AT @ BT #dw/dlr
            m = torch.sum(GT * e) #dloss/dlr

            lr /= down_fac
            while loss1 > loss + c*m*lr:
                lr *= down_fac
                if 2*K < ndim:
                    UT = torch.cat((GT, AT), dim=0).double()
                    V = torch.cat((AT.T, -GT.T), dim=1).double()
                    AT1 = (AT.double() - lr * AT.double() @ V @ torch.pinverse(torch.eye(2*K, dtype=torch.double, device=x.device)+lr/2*UT@V) @ UT).to(torch.get_default_dtype())
                else:
                    AT1 = (AT.double() @ (torch.eye(ndim, dtype=torch.double, device=x.device)-lr/2*BT.double()) @ torch.pinverse(torch.eye(ndim, dtype=torch.double, device=x.device)+lr/2*BT.double())).to(torch.get_default_dtype())
            
                if x2 == 'gaussian':
                    loss1 = -ObjectiveG(AT1 @ x.T, pg, p, w=weight)
                else:
                    loss1 = -Objective(AT1 @ x.T, AT1 @ x2.T, p)
        
            if torch.max(torch.abs(AT1-AT)) < eps:
                AT = AT1
                break
        
            lr *= up_fac
            AT = AT1

    if x2 == 'gaussian':
        WD = ObjectiveG(AT @ x.T, pg, p, w=weight, perdim=False)
    else:
        WD = Objective(AT @ x.T, AT @ x2.T, p, perdim=False)
    return AT.T, WD**(1/p)
//  [python]
def _orthogonalize(matrix, epsilon=0):
    """
    Decide between Gram-Schmidt or QR factorization to orthogonalize the matrix.
    QR factorization doesn't work with half-precision, but it is usually faster with a rank > 2.
    """
    assert len(matrix.shape) == 2 and matrix.shape[1] <= matrix.shape[0]

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

    rank = matrix.shape[1]
    dtype = matrix.dtype
    if rank <= 2 or dtype in [torch.float16, torch.bfloat16]:
        _orthogonalize_gram_schmidt(matrix, epsilon=epsilon)
    else:
        torch.linalg.qr(
            matrix,
            out=(
                matrix,
                torch.empty(rank, rank, device=matrix.device, dtype=dtype)
            )
        )
//  [python]
def project_parameters(self):
        self.U.weight.data = torch.qr(self.U.weight.T)[0].T
//  [python]
def project_parameters(self):
        self.U.weight.data = torch.linalg.qr(self.U.weight.T)[0].T
//  [python]
def test_qr_vs_numpy(self, device, dtype):
        """
        test torch.linalg.qr vs numpy.linalg.qr
        """
        sizes_to_test = [
            (7, 5),
            (5, 7),
            (5, 0),    # empty
            (0, 5),    # empty
        ]
        for size in sizes_to_test:
            t = torch.randn(size, device=device, dtype=dtype)
            np_t = t.cpu().numpy()
            for mode in ['reduced', 'complete']:
                exp_q, exp_r = np.linalg.qr(np_t, mode=mode)
                q, r = t.qr(mode=mode)
                self.assertEqual(q, exp_q)
                self.assertEqual(r, exp_r)
            #
            # for mode='r' we need a special logic because numpy returns only r
            exp_r = np.linalg.qr(np_t, mode='r')
            q, r = t.qr(mode='r')
            # check that q is empty
            self.assertEqual(q.shape, (0,))
            self.assertEqual(q.dtype, t.dtype)
            self.assertEqual(q.device, t.device)
            # check r
            self.assertEqual(r, exp_r)
//  [python]
def test_qr_vs_numpy(self, device, dtype):
        """
        test torch.linalg.qr vs numpy.linalg.qr
        """
        sizes_to_test = [
            (7, 5),
            (5, 7),
            (5, 0),    # empty
            (0, 5),    # empty
        ]
        for size in sizes_to_test:
            t = torch.randn(size, device=device, dtype=dtype)
            np_t = t.cpu().numpy()
            for mode in ['reduced', 'complete']:
                exp_q, exp_r = np.linalg.qr(np_t, mode=mode)
                q, r = torch.linalg.qr(t, mode=mode)
                self.assertEqual(q, exp_q)
                self.assertEqual(r, exp_r)
            #
            # for mode='r' we need a special logic because numpy returns only r
            exp_r = np.linalg.qr(np_t, mode='r')
            q, r = torch.linalg.qr(t, mode='r')
            # check that q is empty
            self.assertEqual(q.shape, (0,))
            self.assertEqual(q.dtype, t.dtype)
            self.assertEqual(q.device, t.device)
            # check r
            self.assertEqual(r, exp_r)
//  [python]
def test_linalg_qr_autograd_errors(self, device, dtype):
        # torch.linalg.qr(mode='r') returns only 'r' and discards 'q', but
        # without 'q' you cannot compute the backward pass. Check that
        # linalg_qr_backward complains cleanly in that case.
        inp = torch.randn((5, 7), device=device, dtype=dtype, requires_grad=True)
        q, r = torch.qr(inp, mode='r')
        self.assertEqual(q.shape, (0,))  # empty tensor
        b = torch.sum(r)
        with self.assertRaisesRegex(RuntimeError,
                                    "The derivative of qr is not implemented when mode='r'"):
            b.backward()
        #
        inp = torch.randn((7, 5), device=device, dtype=dtype, requires_grad=True)
        q, r = torch.qr(inp, mode='complete')
        b = torch.sum(r)
        with self.assertRaisesRegex(RuntimeError,
                                    "The derivative of qr is not implemented when mode='complete' and nrows > ncols"):
            b.backward()
//  [python]
def test_linalg_qr_autograd_errors(self, device, dtype):
        # torch.linalg.qr(mode='r') returns only 'r' and discards 'q', but
        # without 'q' you cannot compute the backward pass. Check that
        # linalg_qr_backward complains cleanly in that case.
        inp = torch.randn((5, 7), device=device, dtype=dtype, requires_grad=True)
        q, r = torch.linalg.qr(inp, mode='r')
        self.assertEqual(q.shape, (0,))  # empty tensor
        b = torch.sum(r)
        with self.assertRaisesRegex(RuntimeError,
                                    "The derivative of qr is not implemented when mode='r'"):
            b.backward()
        #
        inp = torch.randn((7, 5), device=device, dtype=dtype, requires_grad=True)
        q, r = torch.linalg.qr(inp, mode='complete')
        b = torch.sum(r)
        with self.assertRaisesRegex(RuntimeError,
                                    "The derivative of qr is not implemented when mode='complete' and nrows > ncols"):
            b.backward()
//  [python]
def test_ormqr(self, device, dtype):

        def run_test(batch, m, n, fortran_contiguous):
            A = make_tensor((*batch, m, n), dtype=dtype, device=device)
            reflectors, tau = torch.geqrf(A)
            if not fortran_contiguous:
                self.assertTrue(reflectors.mT.is_contiguous())
                reflectors = reflectors.contiguous()

            # Q is of size m x m
            Q, _ = torch.qr(A, some=True)
            C_right = make_tensor((*batch, m, n), dtype=dtype, device=device)
            C_left = make_tensor((*batch, n, m), dtype=dtype, device=device)

            expected = Q @ C_right
            with pytorch_op_timer():
                actual = torch.ormqr(reflectors, tau, C_right, left=True, transpose=False)
            self.assertEqual(expected, actual)

            expected = C_left @ Q
            with pytorch_op_timer():
                actual = torch.ormqr(reflectors, tau, C_left, left=False, transpose=False)
            self.assertEqual(expected, actual)

            expected = Q.mH @ C_right
            with pytorch_op_timer():
                actual = torch.ormqr(reflectors, tau, C_right, left=True, transpose=True)
            self.assertEqual(expected, actual)

            expected = C_left @ Q.mH
            with pytorch_op_timer():
                actual = torch.ormqr(reflectors, tau, C_left, left=False, transpose=True)
            self.assertEqual(expected, actual)

            # if tau is all zeros then the implicit matrix Q is the identity matrix
            # so the actual result should be C_right in this case
            zero_tau = torch.zeros_like(tau)
            with pytorch_op_timer():
                actual = torch.ormqr(reflectors, zero_tau, C_right, left=True, transpose=False)
            self.assertEqual(C_right, actual)

        batches = [(), (0, ), (2, ), (2, 1)]
        ns = [5, 2, 0]
        for batch, (m, n), fortran_contiguous in product(batches, product(ns, ns), [True, False]):
            run_test(batch, m, n, fortran_contiguous)
//  [python]
def test_ormqr(self, device, dtype):

        def run_test(batch, m, n, fortran_contiguous):
            A = make_tensor((*batch, m, n), dtype=dtype, device=device)
            reflectors, tau = torch.geqrf(A)
            if not fortran_contiguous:
                self.assertTrue(reflectors.mT.is_contiguous())
                reflectors = reflectors.contiguous()

            # Q is of size m x m
            Q, _ = torch.linalg.qr(A, mode='complete')
            C_right = make_tensor((*batch, m, n), dtype=dtype, device=device)
            C_left = make_tensor((*batch, n, m), dtype=dtype, device=device)

            expected = Q @ C_right
            with pytorch_op_timer():
                actual = torch.ormqr(reflectors, tau, C_right, left=True, transpose=False)
            self.assertEqual(expected, actual)

            expected = C_left @ Q
            with pytorch_op_timer():
                actual = torch.ormqr(reflectors, tau, C_left, left=False, transpose=False)
            self.assertEqual(expected, actual)

            expected = Q.mH @ C_right
            with pytorch_op_timer():
                actual = torch.ormqr(reflectors, tau, C_right, left=True, transpose=True)
            self.assertEqual(expected, actual)

            expected = C_left @ Q.mH
            with pytorch_op_timer():
                actual = torch.ormqr(reflectors, tau, C_left, left=False, transpose=True)
            self.assertEqual(expected, actual)

            # if tau is all zeros then the implicit matrix Q is the identity matrix
            # so the actual result should be C_right in this case
            zero_tau = torch.zeros_like(tau)
            with pytorch_op_timer():
                actual = torch.ormqr(reflectors, zero_tau, C_right, left=True, transpose=False)
            self.assertEqual(C_right, actual)

        batches = [(), (0, ), (2, ), (2, 1)]
        ns = [5, 2, 0]
        for batch, (m, n), fortran_contiguous in product(batches, product(ns, ns), [True, False]):
            run_test(batch, m, n, fortran_contiguous)
//  [python]
def test_qr_error_cases(self, device, dtype):
        t1 = torch.randn(5, device=device, dtype=dtype)
        with self.assertRaisesRegex(RuntimeError, 'qr input should have at least 2 dimensions, but has 1 dimensions instead'):
            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, 'qr input should have at least 2 dimensions, but has 1 dimensions instead'):
            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 test_qr_out(self, device, dtype):
        """
        test torch.linalg.qr(out=...) vs torch.lingalg.qr
        """
        sizes_to_test = [
            (7, 5),
            (5, 7),
            (5, 0),    # empty
            (0, 5),    # empty
        ]
        for size in sizes_to_test:
            t = torch.randn(size, device=device, dtype=dtype)
            np_t = t.cpu().numpy()
            for mode in ['reduced', 'complete', 'r']:
                q, r = torch.linalg.qr(t, mode=mode)
                out = (torch.empty((0), dtype=dtype, device=device),
                       torch.empty((0), dtype=dtype, device=device))
                q2, r2 = torch.linalg.qr(t, mode=mode, out=out)
                self.assertIs(q2, out[0])
                self.assertIs(r2, out[1])
                self.assertEqual(q2, q)
                self.assertEqual(r2, r)
//  [python]
def orthogonal_matrix_chunk(cols, device=None):
    unstructured_block = torch.randn((cols, cols), device=device)
    if TORCH_GE_1_8_0:
        q, r = torch.linalg.qr(unstructured_block.cpu(), mode='reduced')
    else:
        q, r = torch.qr(unstructured_block.cpu(), some=True)
    q, r = map(lambda t: t.to(device), (q, r))
    return q.t()
//  [python]
def test_qr_batched(self, device, dtype):
        """
        test torch.linalg.qr vs numpy.linalg.qr. We need some special logic
        because numpy does not support batched qr
        """
        def np_qr_batched(a, mode):
            """poor's man batched version of np.linalg.qr"""
            all_q = []
            all_r = []
            for matrix in a:
                result = np.linalg.qr(matrix, mode=mode)
                if mode == 'r':
                    all_r.append(result)
                else:
                    q, r = result
                    all_q.append(q)
                    all_r.append(r)
            if mode == 'r':
                return np.array(all_r)
            else:
                return np.array(all_q), np.array(all_r)

        t = torch.randn((3, 7, 5), device=device, dtype=dtype)
        np_t = t.cpu().numpy()
        for mode in ['reduced', 'complete']:
            exp_q, exp_r = np_qr_batched(np_t, mode=mode)
            q, r = t.qr(mode=mode)
            self.assertEqual(q, exp_q)
            self.assertEqual(r, exp_r)
        # for mode='r' we need a special logic because numpy returns only r
        exp_r = np_qr_batched(np_t, mode='r')
        q, r = t.qr(mode='r')
        # check that q is empty
        self.assertEqual(q.shape, (0,))
        self.assertEqual(q.dtype, t.dtype)
        self.assertEqual(q.device, t.device)
        # check r
        self.assertEqual(r, exp_r)
//  [python]
def test_qr_batched(self, device, dtype):
        """
        test torch.linalg.qr vs numpy.linalg.qr. We need some special logic
        because numpy does not support batched qr
        """
        def np_qr_batched(a, mode):
            """poor's man batched version of np.linalg.qr"""
            all_q = []
            all_r = []
            for matrix in a:
                result = np.linalg.qr(matrix, mode=mode)
                if mode == 'r':
                    all_r.append(result)
                else:
                    q, r = result
                    all_q.append(q)
                    all_r.append(r)
            if mode == 'r':
                return np.array(all_r)
            else:
                return np.array(all_q), np.array(all_r)

        t = torch.randn((3, 7, 5), device=device, dtype=dtype)
        np_t = t.cpu().numpy()
        for mode in ['reduced', 'complete']:
            exp_q, exp_r = np_qr_batched(np_t, mode=mode)
            q, r = torch.linalg.qr(t, mode=mode)
            self.assertEqual(q, exp_q)
            self.assertEqual(r, exp_r)
        # for mode='r' we need a special logic because numpy returns only r
        exp_r = np_qr_batched(np_t, mode='r')
        q, r = torch.linalg.qr(t, mode='r')
        # check that q is empty
        self.assertEqual(q.shape, (0,))
        self.assertEqual(q.dtype, t.dtype)
        self.assertEqual(q.device, t.device)
        # check r
        self.assertEqual(r, exp_r)
//  [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 = 'reduced')
    q, r = map(lambda t: t.to(device), (q, r))
    return q.t()
//  [python]
def basis(A):
    """Return orthogonal basis of A columns.
    """
    if A.is_cuda:
        # torch.orgqr is not available in CUDA
        Q, R = torch.qr(A)
    else:
        Q, R = torch.geqrf(A)
        Q = torch.orgqr(Q, R)
    return Q
//  [python]
def basis(A):
    """Return orthogonal basis of A columns.
    """
    if A.is_cuda:
        # torch.orgqr is not available in CUDA
        Q = torch.linalg.qr(A).Q
    else:
        Q = torch.orgqr(*torch.geqrf(A))
    return Q
//  [python]
def orthogonalize(matrix, overwrite=False):
    """Orthogonalization of given matrix.

    :param matrix: matrix to orthogonalize, needs to be compatible with either
      ``scipy.linalg.qr`` or ``torch.linalg.qr``.
    :param overwrite: If true, ``matrix[:] = Q`` will also be performed.
    :returns: Orthogonal matrix ``Q`` such that ``matrix = Q @ R`` as per the
      QR decomposition.
    """
    h, w = matrix.shape
    assert h >= w, "Only non-fat matrices supported!"
    #
    if isinstance(matrix, torch.Tensor):
        Q = torch.qr(matrix, some=True)[0]
    else:
        Q = scipy.linalg.qr(
            matrix,
            mode="economic",
            pivoting=True,
        )[0]
    #
    if overwrite:
        matrix[:] = Q
    return Q
//  [python]
def orthogonalize(matrix, overwrite=False):
    """Orthogonalization of given matrix.

    :param matrix: matrix to orthogonalize, needs to be compatible with either
      ``scipy.linalg.qr`` or ``torch.linalg.qr``.
    :param overwrite: If true, ``matrix[:] = Q`` will also be performed.
    :returns: Orthogonal matrix ``Q`` such that ``matrix = Q @ R`` as per the
      QR decomposition.
    """
    h, w = matrix.shape
    assert h >= w, "Only non-fat matrices supported!"
    #
    if isinstance(matrix, torch.Tensor):
        Q = torch.linalg.qr(matrix, mode="reduced")[0]
    else:
        Q = scipy.linalg.qr(
            matrix,
            mode="economic",
            pivoting=True,
        )[0]
    #
    if overwrite:
        matrix[:] = Q
    return Q
//  [python]
def test_gvp_no_vector_gate_equivariance(self):
        torch.manual_seed(0)
        in_dims = (3, 4)
        out_dims = (5, 6)
        vector_gate = False
        gvp = GVP(in_dims, out_dims, vector_gate=vector_gate).double()
        s = torch.randn(10, 3, dtype=torch.float64)
        V = torch.randn(10, 4, 3, dtype=torch.float64)
        x1 = (s, V)
        R, _ = torch.qr(torch.randn(3, 3, dtype=torch.float64))
        x2 = (s, V @ R)
        s_dash1, V_dash1 = gvp(x1)
        s_dash2, V_dash2 = gvp(x2)

        self.assertEqual(s_dash1.shape, (10, 5))
        self.assertEqual(V_dash1.shape, (10, 6, 3))
        self.assertEqual(s_dash2.shape, (10, 5))
        self.assertEqual(V_dash2.shape, (10, 6, 3))
        self.assertTrue(torch.allclose(s_dash1, s_dash2))
        self.assertTrue(torch.allclose(V_dash1 @ R, V_dash2))
//  [python]
def test_gvp_no_vector_gate_equivariance(self):
        torch.manual_seed(0)
        in_dims = (3, 4)
        out_dims = (5, 6)
        vector_gate = False
        gvp = GVP(in_dims, out_dims, vector_gate=vector_gate).double()
        s = torch.randn(10, 3, dtype=torch.float64)
        V = torch.randn(10, 4, 3, dtype=torch.float64)
        x1 = (s, V)
        R, _ = torch.linalg.qr(torch.randn(3, 3, dtype=torch.float64))
        x2 = (s, V @ R)
        s_dash1, V_dash1 = gvp(x1)
        s_dash2, V_dash2 = gvp(x2)

        self.assertEqual(s_dash1.shape, (10, 5))
        self.assertEqual(V_dash1.shape, (10, 6, 3))
        self.assertEqual(s_dash2.shape, (10, 5))
        self.assertEqual(V_dash2.shape, (10, 6, 3))
        self.assertTrue(torch.allclose(s_dash1, s_dash2))
        self.assertTrue(torch.allclose(V_dash1 @ R, V_dash2))
//  [python]
def test_gvp_equivariance(self):
        torch.manual_seed(0)
        in_dims = (3, 4)
        out_dims = (5, 6)
        vector_gate = True
        gvp = GVP(in_dims, out_dims, vector_gate=vector_gate).double()
        s = torch.randn(10, 3, dtype=torch.float64)
        V = torch.randn(10, 4, 3, dtype=torch.float64)
        x1 = (s, V)
        R, _ = torch.qr(torch.randn(3, 3, dtype=torch.float64))
        x2 = (s, V @ R)
        s_dash1, V_dash1 = gvp(x1)
        s_dash2, V_dash2 = gvp(x2)

        self.assertEqual(s_dash1.shape, (10, 5))
        self.assertEqual(V_dash1.shape, (10, 6, 3))
        self.assertEqual(s_dash2.shape, (10, 5))
        self.assertEqual(V_dash2.shape, (10, 6, 3))
        self.assertTrue(torch.allclose(s_dash1, s_dash2))
        self.assertTrue(torch.allclose(V_dash1 @ R, V_dash2))
//  [python]
def test_gvp_equivariance(self):
        torch.manual_seed(0)
        in_dims = (3, 4)
        out_dims = (5, 6)
        vector_gate = True
        gvp = GVP(in_dims, out_dims, vector_gate=vector_gate).double()
        s = torch.randn(10, 3, dtype=torch.float64)
        V = torch.randn(10, 4, 3, dtype=torch.float64)
        x1 = (s, V)
        R, _ = torch.linalg.qr(torch.randn(3, 3, dtype=torch.float64))
        x2 = (s, V @ R)
        s_dash1, V_dash1 = gvp(x1)
        s_dash2, V_dash2 = gvp(x2)

        self.assertEqual(s_dash1.shape, (10, 5))
        self.assertEqual(V_dash1.shape, (10, 6, 3))
        self.assertEqual(s_dash2.shape, (10, 5))
        self.assertEqual(V_dash2.shape, (10, 6, 3))
        self.assertTrue(torch.allclose(s_dash1, s_dash2))
        self.assertTrue(torch.allclose(V_dash1 @ R, V_dash2))
//  [python]
def orthogonal_matrix_chunk(cols, device = None):
    unstructured_block = torch.randn((cols, cols), device = device)
    if TORCH_GE_1_8_0:
        q, r = torch.linalg.qr(unstructured_block.cpu(), mode = 'reduced')
    else:
        q, r = torch.qr(unstructured_block.cpu(), some = True)
    q, r = map(lambda t: t.to(device), (q, r))
    return q.t()
//  [python]
def hutchpp(A: lo.LinearOperator, m: int=102, *, _gen_data_on_device=None):
    n, ell = A.shape
    assert n == ell
    assert m % 3 == 0
    if n <= m:
        return exact_trace(A)
    k = m // 3
    S = _sample_dist('rademacher', (n, k), _gen_data_on_device, A.device)
    G = _sample_dist('rademacher', (n, k), _gen_data_on_device, A.device)

    BS = lo.operator_matrix_product(A, S)
    Q, _ = torch.qr(BS)
    G_prime = G - Q @ (Q.T @ G)
    exact = 0
    for i in range(Q.shape[1]):
        z = Q[:,i]
        exact = exact + (z * (A @ z)).sum()
    
    hutchinson = torch.empty(k, device=A.device)
    for i in range(k):
        z = G_prime[:,i]
        hutchinson[i] = (z * (A @ z)).sum()

    return exact + torch.mean(hutchinson), torch.std(hutchinson) / math.sqrt(m)
//  [python]
def hutchpp(A: lo.LinearOperator, m: int=102, *, _gen_data_on_device=None):
    n, ell = A.shape
    assert n == ell
    assert m % 3 == 0
    if n <= m:
        return exact_trace(A)
    k = m // 3
    S = _sample_dist('rademacher', (n, k), _gen_data_on_device, A.device)
    G = _sample_dist('rademacher', (n, k), _gen_data_on_device, A.device)

    BS = lo.operator_matrix_product(A, S)
    Q, _ = torch.linalg.qr(BS)
    G_prime = G - Q @ (Q.T @ G)
    exact = 0
    for i in range(Q.shape[1]):
        z = Q[:,i]
        exact = exact + (z * (A @ z)).sum()
    
    hutchinson = torch.empty(k, device=A.device)
    for i in range(k):
        z = G_prime[:,i]
        hutchinson[i] = (z * (A @ z)).sum()

    return exact + torch.mean(hutchinson), torch.std(hutchinson) / math.sqrt(m)
//  [python]
def orthogonalize(vs):  # vs ~ [batch, n, N]
    q, r = torch.qr(vs.float().transpose(-1, -2))
    return q.transpose(-1, -2)
//  [python]
def orthogonalize(vs):  # vs ~ [batch, n, N]
    q = torch.linalg.qr(vs.float().transpose(-1, -2), 'reduced')[0]
    return q.transpose(-1, -2)
//  [python]
def generate_ps(self, inp, N, level=None):
        C, H, W = inp.shape
        h_use, w_use = int(H / self.factor), int(W / self.factor)
        p_signal = torch.zeros((N,) + inp.shape, device=inp.device)
        rv = torch.randn((N, C, h_use, w_use))
        rv_ortho, _ = torch.qr(rv)
        rv_ortho = rv_ortho.to(inp.device)
        p_signal[:, :, :h_use, :w_use] = rv_ortho
        im_blocks = blockify(p_signal, H)
        dct_blocks = block_idct(im_blocks)
        ps = deblockify(dct_blocks, (H, W))
        return ps
//  [python]
def generate_ps(self, inp, N, level=None):
        C, H, W = inp.shape
        h_use, w_use = int(H / self.factor), int(W / self.factor)
        p_signal = torch.zeros((N,) + inp.shape, device=inp.device)
        # rv = torch.randn((N, C, h_use, w_use), device=inp.device)
        # rv_ortho, _ = torch.linalg.qr(rv, mode='complete')
        rv = torch.randn((N, C, h_use, w_use))
        rv_ortho, _ = torch.linalg.qr(rv, mode="complete")
        rv_ortho = rv_ortho.to(inp.device)
        p_signal[:, :, :h_use, :w_use] = rv_ortho
        im_blocks = blockify(p_signal, H)
        # Orthonormal by default
        dct_blocks = block_idct(im_blocks)
        ps = deblockify(dct_blocks, (H, W))
        return ps
//  [python]
def project(y, X):
    y = y.to(torch.float)
    X = (X.T).to(torch.float)  # make sure basis is columns

    Q, R = torch.qr(X)
    beta = torch.triangular_solve(Q.T @ y.unsqueeze(-1), R).solution
    proj = (X @ beta)
    return proj.squeeze(-1)
//  [python]
def project(y, X):
    y = y.to(torch.float)
    X = (X.T).to(torch.float)  # make sure basis is columns

    Q, R = torch.linalg.qr(X)
    beta = torch.triangular_solve(Q.T @ y.unsqueeze(-1), R).solution
    proj = (X @ beta)
    return proj.squeeze(-1)
//  [python]
def project_points_to_planes(y, X):
    y = y.to(torch.float).unsqueeze(1)
    X = (X.transpose(-1, -2)).to(torch.float)  # make sure basis is the columns

    Q, R = torch.qr(X)
    Q = Q.unsqueeze(1)

    # solve A * beta = b
    b = Q.transpose(-1, -2) @ y.unsqueeze(-1)
    b.squeeze_(1)

    beta = torch.triangular_solve(b, R).solution
    proj = (X @ beta)
    return proj.squeeze(-1)
//  [python]
def project_points_to_planes(y, X):
    y = y.to(torch.float).unsqueeze(1)
    X = (X.transpose(-1, -2)).to(torch.float)  # make sure basis is the columns

    Q, R = torch.linalg.qr(X)
    Q = Q.unsqueeze(1)

    # solve A * beta = b
    b = Q.transpose(-1, -2) @ y.unsqueeze(-1)
    b.squeeze_(1)

    beta = torch.triangular_solve(b, R).solution
    proj = (X @ beta)
    return proj.squeeze(-1)
//  [python]
def xtrace(A, m: int=80, *, _gen_data_on_device=None):
    n, n1 = A.shape
    assert n == n1
    if n <= m:
        return exact_trace(A)

    m = m // 2
    def normalize_columns(M):
        return M / torch.linalg.vector_norm(M, dim=0)

    def diag_of_AB(A, B):
        return torch.sum(A * B, dim=0)

    Z = _sample_dist('gaussian', (n, m), _gen_data_on_device, A.device)
    Omega = math.sqrt(n) * normalize_columns(Z)
    Y = A @ Omega
    Q, R = torch.qr(Y)

    W = Q.T @ Omega
    S = normalize_columns(torch.linalg.inv(R).T)
    scale = (n - m + 1) / (n - torch.linalg.vector_norm(W, dim=0)**2 +
                           torch.abs(diag_of_AB(S, W) * torch.linalg.vector_norm(S, dim=0))**2)
    Z = A @ Q
    H = Q.T @ Z
    HW = H @ W
    T = Z.T @ Omega
    dSW = diag_of_AB(S, W)
    dSHS = diag_of_AB(S, H @ S)
    dTW = diag_of_AB(T, W)
    dWHW = diag_of_AB(W, HW)
    dSRmHW = diag_of_AB(S, R - HW)
    dTmHRS = diag_of_AB(T - H.T @ W, S)

    ests = torch.trace(H) - dSHS + (
            dWHW - dTW + dTmHRS * dSW + torch.abs(dSW)**2 * dSHS + dSW * dSRmHW) * scale

    return torch.mean(ests), torch.std(ests) / math.sqrt(m)
//  [python]
def xtrace(A, m: int=80, *, _gen_data_on_device=None):
    n, n1 = A.shape
    assert n == n1
    if n <= m:
        return exact_trace(A)

    m = m // 2
    def normalize_columns(M):
        return M / torch.linalg.vector_norm(M, dim=0)

    def diag_of_AB(A, B):
        return torch.sum(A * B, dim=0)

    Z = _sample_dist('gaussian', (n, m), _gen_data_on_device, A.device)
    Omega = math.sqrt(n) * normalize_columns(Z)
    Y = A @ Omega
    Q, R = torch.linalg.qr(Y)

    W = Q.T @ Omega
    S = normalize_columns(torch.linalg.inv(R).T)
    scale = (n - m + 1) / (n - torch.linalg.vector_norm(W, dim=0)**2 +
                           torch.abs(diag_of_AB(S, W) * torch.linalg.vector_norm(S, dim=0))**2)
    Z = A @ Q
    H = Q.T @ Z
    HW = H @ W
    T = Z.T @ Omega
    dSW = diag_of_AB(S, W)
    dSHS = diag_of_AB(S, H @ S)
    dTW = diag_of_AB(T, W)
    dWHW = diag_of_AB(W, HW)
    dSRmHW = diag_of_AB(S, R - HW)
    dTmHRS = diag_of_AB(T - H.T @ W, S)

    ests = torch.trace(H) - dSHS + (
            dWHW - dTW + dTmHRS * dSW + torch.abs(dSW)**2 * dSHS + dSW * dSRmHW) * scale

    return torch.mean(ests), torch.std(ests) / math.sqrt(m)
//  [python]
def compute_orth_rep(self, rep):
        """
        Use QR decomposition to orthogonalise the representation but retain the subspace spanned by the columns. 

        Args:
            rep (torch.tensor): (group.order, dim^2) tensor of representation

        Returns:
            torch.tensor: (group.order, dim^2) tensor with orthonormal columns en
        """
        
        orth_rep = rep.reshape(self.order, self.dim * self.dim)

        # TODO: this is a hack to not double count in the cyclic group
        if self.group_acronym == 'C' and self.dim == 2:
            orth_rep = orth_rep[:, :2] 
        q, r = torch.qr(orth_rep)
        orth_rep = q
        
        return orth_rep
//  [python]
def compute_orth_rep(self, rep):
        """
        Use QR decomposition to orthogonalise the representation but retain the subspace spanned by the columns. 

        Args:
            rep (torch.tensor): (group.order, dim^2) tensor of representation

        Returns:
            torch.tensor: (group.order, dim^2) tensor with orthonormal columns en
        """
        
        orth_rep = rep.reshape(self.order, self.dim * self.dim)

        # TODO: this is a hack to not double count in the cyclic group
        if self.group_acronym == 'C' and self.dim == 2:
            orth_rep = orth_rep[:, :2] 
        orth_rep = torch.linalg.qr(orth_rep)[0]
        
        return orth_rep
//  [python]
def xnystrace(A: lo.LinearOperator, m: int=80, *, _gen_data_on_device=None):
    n, n1 = A.shape
    assert n == n1
    if n <= m:
        return exact_trace(A)
    m = m // 2
    def normalize_columns(M):
        return M / torch.linalg.vector_norm(M, dim=0)

    def diag_of_AB(A, B):
        return torch.sum(A * B, dim=0)

    Z = _sample_dist('gaussian', (n, m), _gen_data_on_device, A.device)
    Omega = math.sqrt(n) * normalize_columns(Z)
    Y = A @ Omega

    nu = torch.finfo(Omega.dtype).eps * torch.linalg.vector_norm(Y) / math.sqrt(n)
    Y = Y + nu * Omega


    Q, R = torch.qr(Y)
    H = Omega.T @ Y
    C = torch.linalg.cholesky((H + H.T) / 2, upper=True)

    B = torch.linalg.solve_triangular(C, R, upper=True, left=False) # Double check this

    QQ, RR = torch.qr(Omega)
    WW = QQ.T @ Omega
    SS = normalize_columns(torch.linalg.inv(RR).T)

    scale = (n - m + 1) / (n - torch.linalg.vector_norm(WW, dim=0)**2 +
                           torch.abs(diag_of_AB(SS, WW) * torch.linalg.vector_norm(SS, dim=0))**2)
    W = Q.T @ Omega
    S =  torch.linalg.solve_triangular(C.T, B, upper=False, left=False) * torch.diag(torch.linalg.inv(H))**(-0.5) # Double check triangular solve
    dSW = diag_of_AB(S, W)
    ests = torch.linalg.norm(B, 'fro')**2 - torch.linalg.vector_norm(S, dim=0)**2 + torch.abs(dSW)**2 * scale - nu * n
    return torch.mean(ests), torch.std(ests) / math.sqrt(m)