jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def orthogonal_(
    tensor,
    gain=1,
    generator: _Optional[torch.Generator] = None,
):
    r"""Fill the input `Tensor` with a (semi) orthogonal matrix.

    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
        generator: the torch Generator to sample from (default: None)

    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=generator)

    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,
    generator: _Optional[torch.Generator] = None,
):
    r"""Fill the input `Tensor` with a (semi) orthogonal matrix.

    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
        generator: the torch Generator to sample from (default: None)

    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=generator)

    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 forward(self, input):
        # input: (bs*t) x 512

        weight = self.weight + 1e-8
        Q, R = torch.qr(weight)  # get eignvector, orthogonal [n1, n2, n3, n4]
        #Q, R = gram_schmidt_qr(weight) # custom OP

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

            return out
//  [python]
def forward(self, input):
        # input: (bs*t) x 512

        weight = self.weight + 1e-8
        Q, R = torch.linalg.qr(weight, 'reduced')  # get eignvector, orthogonal [n1, n2, n3, n4]
        #Q, R = gram_schmidt_qr(weight) # custom OP

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

            return out
//  [python]
def make_orthogonal(tensor):
    rows = tensor.size(0)
    cols = tensor.numel() // rows
    flattened = torch.reshape(tensor, (rows, cols))
    if rows < cols:
        flattened.t_()
    q, r = torch.qr(flattened)
    d = torch.diag(r, 0)
    ph = d.sign()
    q *= ph
    if rows < cols:
        q.t_()        
    return q
//  [python]
def make_orthogonal(tensor):
    #similar to torch.nn.init.orthogonal_
    rows = tensor.size(0)
    cols = tensor.numel() // rows
    flattened = torch.reshape(tensor, (rows, cols))
    if rows < cols:
        flattened.t_()
    q, r = torch.linalg.qr(flattened)
    d = torch.diag(r, 0)
    ph = d.sign()
    q *= ph
    if rows < cols:
        q.t_()        
    return q
//  [python]
def ttd_tc(x, chi=-1, svd=True):
    """
    :param x: tensor to be decomposed
    :param chi: dimension cut-off. Use QR decomposition when dc=None;
                use SVD but don'tc truncate when dc=-1
    :param svd: use svd or qr
    :return tensors: tensors in the TT form
    :return lm: singular values in each decomposition (calculated when dc is not None)
    """
    dims = x.shape
    ndim = x.ndimension()
    dimL = 1
    tensors = list()
    lm = list()
    if chi is None:
        chi = -1
    for n in range(ndim-1):
        if (chi < 0) and (not svd):  # No truncation
            q, x = tc.qr(x.reshape(dimL*dims[n], -1))
            dimL1 = x.shape[0]
        else:
            q, s, v = tc.linalg.svd(x.reshape(dimL*dims[n], -1))
            if chi > 0:
                dc = min(chi, s.numel())
            else:
                dc = s.numel()
            q = q[:, :dc]
            s = s[:dc].to(q.dtype)
            lm.append(s)
            x = tc.diag(s).mm(v[:dc, :])
            dimL1 = dc
        tensors.append(q.reshape(dimL, dims[n], dimL1))
        dimL = dimL1
    tensors.append(x.reshape(dimL, dims[-1], 1))
    # tensors[0] = tensors[0][0, :, :]
    return tensors, lm
//  [python]
def ttd_tc(x, chi=-1, svd=True):
    """
    :param x: tensor to be decomposed
    :param chi: dimension cut-off. Use QR decomposition when dc=None;
                use SVD but don'tc truncate when dc=-1
    :param svd: use svd or qr
    :return tensors: tensors in the TT form
    :return lm: singular values in each decomposition (calculated when dc is not None)
    """
    dims = x.shape
    ndim = x.ndimension()
    dimL = 1
    tensors = list()
    lm = list()
    if chi is None:
        chi = -1
    for n in range(ndim-1):
        if (chi < 0) and (not svd):  # No truncation
            q, x = tc.linalg.qr(x.reshape(dimL*dims[n], -1))
            dimL1 = x.shape[0]
        else:
            q, s, v = tc.linalg.svd(x.reshape(dimL*dims[n], -1))
            if chi > 0:
                dc = min(chi, s.numel())
            else:
                dc = s.numel()
            q = q[:, :dc]
            s = s[:dc].to(q.dtype)
            lm.append(s)
            x = tc.diag(s).mm(v[:dc, :])
            dimL1 = dc
        tensors.append(q.reshape(dimL, dims[n], dimL1))
        dimL = dimL1
    tensors.append(x.reshape(dimL, dims[-1], 1))
    # tensors[0] = tensors[0][0, :, :]
    return tensors, lm
//  [python]
def qr(A):
    try:
        return torch.linalg.qr(A, mode='reduced')
    except AttributeError:
        return torch.qr(A, some=True)
//  [python]
def qr(A):
    try:
        # only for newer pytorch versions
        return torch.linalg.qr(A, mode='reduced')
    except AttributeError:
        return torch.qr(A, some=True)
//  [python]
def get_projected_weight(self, pull, base_weight, weights):
        tr = torch.transpose(base_weight, 0, 1)
        Q, R = torch.qr(tr) # Q is 640x60
        mut = weights @ Q # mut is 5 x 60
        mutnorm = mut / torch.norm(base_weight, dim=1).unsqueeze(0)
        return mutnorm @ base_weight
//  [python]
def get_projected_weight(self, pull, base_weight, weights):
        tr = torch.transpose(base_weight, 0, 1)
        Q, R = torch.linalg.qr(tr) # Q is 640x60
        mut = weights @ Q # mut is 5 x 60
        mutnorm = mut / torch.norm(base_weight, dim=1).unsqueeze(0)
        return mutnorm @ base_weight
//  [python]
def _fit_full(self, x):
        """Fit the model by computing full SVD on X."""
        x = torch.tensor(x,dtype=torch.float32).to(self.device)
        n_samples, n_features = x.shape       
        X_center = x-torch.mean(x, axis=0)   # Center data
        q ,r = torch.qr(X_center)
        U, s, Vt=torch.linalg.svd(r, full_matrices=False)
        self.__Vt=Vt
        explained_variance = (s ** 2) / (n_samples - 1)   # Get variance explained by singular values
        total_var = explained_variance.sum()
        explained_variance_ratio = explained_variance / total_var
        # explained_variance_ratio = list(explained_variance_ratio)
        if self.n_component_ratio>0 and self.n_component_ratio<1:
            ratio_cumsum = stable_cumsum(explained_variance_ratio.cpu().numpy())
            self.n_components = np.searchsorted(ratio_cumsum, self.n_component_ratio, side="right") + 1
            # self.n_components=self.__choose_ratio(explained_r=explained_variance_ratio)  #find how many features we should keep on the compressed rate we select
        elif type(self.n_component_ratio)==int and self.n_component_ratio<n_features:
            self.n_components = self.n_component_ratio
        else:
            AttributeError("n_components error!")
        self.explained_variance_ratio=explained_variance_ratio[:self.n_components]
        return U, s, Vt, q
//  [python]
def _fit_full(self, x):
        """Fit the model by computing full SVD on X."""
        x = torch.tensor(x,dtype=torch.float32).to(self.device)
        n_samples, n_features = x.shape       
        X_center = x-torch.mean(x, axis=0)   # Center data
        q ,r = torch.linalg.qr(X_center)
        U, s, Vt=torch.linalg.svd(r, full_matrices=False)
        self.__Vt=Vt
        explained_variance = (s ** 2) / (n_samples - 1)   # Get variance explained by singular values
        total_var = explained_variance.sum()
        explained_variance_ratio = explained_variance / total_var
        # explained_variance_ratio = list(explained_variance_ratio)
        if self.n_component_ratio>0 and self.n_component_ratio<1:
            ratio_cumsum = stable_cumsum(explained_variance_ratio.cpu().numpy())
            self.n_components = np.searchsorted(ratio_cumsum, self.n_component_ratio, side="right") + 1
            # self.n_components=self.__choose_ratio(explained_r=explained_variance_ratio)  #find how many features we should keep on the compressed rate we select
        elif type(self.n_component_ratio)==int and self.n_component_ratio<n_features:
            self.n_components = self.n_component_ratio
        else:
            AttributeError("n_components error!")
        self.explained_variance_ratio=explained_variance_ratio[:self.n_components]
        return U, s, Vt, q
//  [python]
def transform(self,y):
        """Fit the model with X and apply the dimensionality reduction on y.

        Parameters
        ----------
        X : array-like of shape (n_samples, n_features)
            Training data, where n_samples is the number of samples
            and n_features is the number of features.

        Returns
        -------
        X_new : ndarray of shape (n_samples, n_components)
            Transformed values.

        Notes
        -----
        This method returns a Fortran-ordered array.
        """
        # y = np.array(y)
        y = torch.tensor(y,dtype=torch.float32).to(self.device)
        y_centre = y-torch.mean(y, axis=0)
        q ,r = torch.qr(y_centre)
        y_compress=torch.matmul(r,torch.linalg.inv(self.__Vt)[:,:self.n_components])
        del y_centre,r,y  #del variable and release memory
        y_compressed = torch.matmul(q,y_compress)        
        if "cuda" in str(self.device):
            y_compressed = y_compressed.cpu().numpy() 
        else:
            y_compressed = y_compressed.numpy()
        return y_compressed
//  [python]
def transform(self,y):
        """Fit the model with X and apply the dimensionality reduction on y.

        Parameters
        ----------
        X : array-like of shape (n_samples, n_features)
            Training data, where n_samples is the number of samples
            and n_features is the number of features.

        Returns
        -------
        X_new : ndarray of shape (n_samples, n_components)
            Transformed values.

        Notes
        -----
        This method returns a Fortran-ordered array.
        """
        # y = np.array(y)
        y = torch.tensor(y,dtype=torch.float32).to(self.device)
        y_centre = y-torch.mean(y, axis=0)
        q ,r = torch.linalg.qr(y_centre)
        y_compress=torch.matmul(r,torch.linalg.inv(self.__Vt)[:,:self.n_components])
        del y_centre,r,y  #del variable and release memory
        y_compressed = torch.matmul(q,y_compress)        
        if "cuda" in str(self.device):
            y_compressed = y_compressed.cpu().numpy() 
        else:
            y_compressed = y_compressed.numpy()
        return y_compressed
//  [python]
def _test_linalg_qr(self, A_list):
        for A in A_list:
            Q1, R1 = torch.qr(A)
            Q2, R2 = torch.qr(A.cuda())
            self.assertEqual(Q1, Q2.cpu(), prec=1e-4)
            self.assertEqual(R1, R2.cpu(), prec=1e-3)

            Q1, R1 = torch.qr(A, some=False)
            Q2, R2 = torch.qr(A.cuda(), some=False)
            self.assertEqual(Q1, Q2.cpu(), prec=1e-3)
            self.assertEqual(R1, R2.cpu(), prec=1e-3)

            Q1, R1 = torch.qr(A, some=True)
            Q2, R2 = torch.qr(A.cuda(), some=True)
            self.assertEqual(Q1, Q2.cpu(), prec=1e-3)
            self.assertEqual(R1, R2.cpu(), prec=1e-3)
//  [python]
def _test_linalg_qr(self, A_list):
        for A in A_list:
            Q1, R1 = torch.linalg.qr(A)
            Q2, R2 = torch.linalg.qr(A.cuda())
            self.assertEqual(Q1, Q2.cpu(), prec=1e-4)
            self.assertEqual(R1, R2.cpu(), prec=1e-3)

            Q1, R1 = torch.linalg.qr(A, mode="r")
            Q2, R2 = torch.linalg.qr(A.cuda(), mode="r")
            self.assertEqual(Q1, Q2.cpu(), prec=1e-3)
            self.assertEqual(R1, R2.cpu(), prec=1e-3)

            Q1, R1 = torch.linalg.qr(A, mode="complete")
            Q2, R2 = torch.linalg.qr(A.cuda(), mode="complete")
            self.assertEqual(Q1, Q2.cpu(), prec=1e-3)
            self.assertEqual(R1, R2.cpu(), prec=1e-3)
//  [python]
def right_inverse(self, X, check_in_manifold=True):
        if check_in_manifold and not self.in_manifold(X):
            raise InManifoldError(X, self)
        if self.n != self.k:
            # N will be a completion of X to an orthogonal basis of R^n
            N = X.new_empty(*(self.tensorial_size + (self.n, self.n - self.k)))
            with torch.no_grad():
                N.normal_()
                # We assume for now that X is orthogonal.
                # This will be checked in super().right_inverse()
                # Project N onto the orthogonal complement to X
                # We iterate this twice for this algorithm to be numerically stable
                # This is standard, as done in some stochastic SVD algorithms
                for _ in range(2):
                    N = N - X @ (X.transpose(-2, -1) @ N)
                    # And make it an orthonormal base of the image
                    Q, R = torch.linalg.qr(N)
                    N = Q
                X = torch.cat([X, N], dim=-1)
        return super().right_inverse(X, check_in_manifold=False)[..., : self.k]
//  [python]
def right_inverse(self, X, check_in_manifold=True):
        if check_in_manifold and not self.in_manifold(X):
            raise InManifoldError(X, self)
        if self.n != self.k:
            # N will be a completion of X to an orthogonal basis of R^n
            N = X.new_empty(*(self.tensorial_size + (self.n, self.n - self.k)))
            with torch.no_grad():
                N.normal_()
                # We assume for now that X is orthogonal.
                # This will be checked in super().right_inverse()
                # Project N onto the orthogonal complement to X
                # We iterate this twice for this algorithm to be numerically stable
                # This is standard, as done in some stochastic SVD algorithms
                for _ in range(2):
                    N = N - X @ (X.transpose(-2, -1) @ N)
                    # And make it an orthonormal base of the image
                    N = torch.linalg.qr(N).Q
                X = torch.cat([X, N], dim=-1)
        return super().right_inverse(X, check_in_manifold=False)[..., : self.k]
//  [python]
def __init__(self, K=4, N=4, P=10, noise_power=1, episode_length=6, num_env=4, device=torch.device("cuda:0")):
        self.N = N  # #antennas
        self.K = K  # #users
        self.P = P  # Power
        self.noise_power = noise_power
        self.basis_vectors, _ = torch.qr(torch.rand(2 * self.K * self.N, 2 * self.K * self.N, dtype=torch.float))
        self.subspace_dim = 1
        self.num_env = num_env
        self.device = device
        self.episode_length = episode_length
        self.get_vec_reward = vmap(self.get_reward, in_dims = (0, 0), out_dims = (0))
//  [python]
def __init__(self, K=4, N=4, P=10, noise_power=1, episode_length=6, num_env=4, device=torch.device("cuda:0")):
        self.N = N  # #antennas
        self.K = K  # #users
        self.P = P  # Power
        self.noise_power = noise_power
        self.basis_vectors, _ = torch.linalg.qr(torch.rand(2 * self.K * self.N, 2 * self.K * self.N, dtype=torch.float))
        self.subspace_dim = 1
        self.num_env = num_env
        self.device = device
        self.episode_length = episode_length
        self.get_vec_reward = vmap(self.get_reward, in_dims = (0, 0), out_dims = (0))
//  [python]
def get_fusion_ten(self):
        ''' D4h representation fusion tensor '''
        # construct reflection matrix
        sig_src = torch.block_diag(*[rep_sig[rep] for rep in self.rep_src]) # (d_src, d_src)
        sig_tgt = torch.block_diag(*[rep_sig[rep] for rep in self.rep_tgt]) # (d_tgt, d_tgt)
        sig = torch.kron(sig_src, sig_tgt) # (d_src d_tgt, d_scr d_tgt)
        # construct reflection symmetry projector
        proj = sig - torch.eye(sig.shape[0]).to(sig) # (d_src d_tgt, d_scr d_tgt)
        # find null space
        q, r = torch.qr(proj) # (d_src d_tgt, d_scr d_tgt)
        msk = torch.diag(r) == 0 # (d_scr d_tgt, )
        nullspace = q[:, msk] # (d_src d_tgt, c)
        # construct C4 rotations
        C4_src = torch.block_diag(*[rep_C4[rep] for rep in self.rep_src]) # (d_src, d_src)
        C4_tgt = torch.block_diag(*[rep_C4[rep] for rep in self.rep_tgt]) # (d_tgt, d_tgt)
        C4 = torch.kron(C4_src, C4_tgt) # (d_src d_tgt, d_scr d_tgt)
        fusion_ten = [nullspace]
        for _ in range(3):
            nullspace = C4.mm(nullspace) # (d_src d_tgt, c)
            fusion_ten.append(nullspace)
        fusion_ten = torch.stack(fusion_ten).view(4, self.d_src, self.d_tgt, -1) # (4, d_src, d_tgt, c)
        return fusion_ten
//  [python]
def get_fusion_ten(self):
        ''' D4h representation fusion tensor '''
        # construct reflection matrix
        sig_src = torch.block_diag(*[rep_sig[rep] for rep in self.rep_src]) # (d_src, d_src)
        sig_tgt = torch.block_diag(*[rep_sig[rep] for rep in self.rep_tgt]) # (d_tgt, d_tgt)
        sig = torch.kron(sig_src, sig_tgt) # (d_src d_tgt, d_scr d_tgt)
        # construct reflection symmetry projector
        proj = sig - torch.eye(sig.shape[0]).to(sig) # (d_src d_tgt, d_scr d_tgt)
        # find null space
        q, r = torch.linalg.qr(proj, mode='complete') # (d_src d_tgt, d_scr d_tgt)
        msk = torch.diag(r) == 0 # (d_scr d_tgt, )
        nullspace = q[:, msk] # (d_src d_tgt, c)
        # construct C4 rotations
        C4_src = torch.block_diag(*[rep_C4[rep] for rep in self.rep_src]) # (d_src, d_src)
        C4_tgt = torch.block_diag(*[rep_C4[rep] for rep in self.rep_tgt]) # (d_tgt, d_tgt)
        C4 = torch.kron(C4_src, C4_tgt) # (d_src d_tgt, d_scr d_tgt)
        fusion_ten = [nullspace]
        for _ in range(3):
            nullspace = C4.mm(nullspace) # (d_src d_tgt, c)
            fusion_ten.append(nullspace)
        fusion_ten = torch.stack(fusion_ten).view(4, self.d_src, self.d_tgt, -1) # (4, d_src, d_tgt, c)
        return fusion_ten
//  [python]
def _orthogonal_matrix_chunk(cols, device=None):
    unstructured_block = torch.randn((cols, cols), device=device)
    q, r = torch.qr(unstructured_block.cpu())
    q, r = map(lambda t: t.to(device), (q, r))
    return q.t()
//  [python]
def _orthogonal_matrix_chunk(cols, device=None):
    unstructured_block = torch.randn((cols, cols), device=device)
    q, r = torch.linalg.qr(unstructured_block.cpu(), mode='complete')
    q, r = map(lambda t: t.to(device), (q, r))
    return q.t()
//  [python]
def K_postprocess_step(self):

        for l in self.NN.layer:

            if hasattr(l,'lr') and l.lr:

                if not l.fixed:
                    
                    U_hat = torch.hstack((l.K[:,:l.dynamic_rank],l.U[:,:l.dynamic_rank]))

                    try:
                        U_hat,R = torch.qr(U_hat)
                    except:
                        U_hat,R = np.linalg.qr(U_hat)
                        U_hat = torch.tensor(U_hat)
                    l.U_hat[:,:2*l.dynamic_rank] = U_hat
                    l.M_hat[:2*l.dynamic_rank,:l.dynamic_rank] = l.U_hat[:,:2*l.dynamic_rank].T@l.U[:,:l.dynamic_rank]
                
                else:

                    try:
                        U_hat,R = torch.qr(l.K)

                    except:
                        U_hat,R = np.linalg.qr(U_hat)
                        U_hat = torch.tensor(U_hat)
                    l.M_hat.data = U_hat.T@l.U.data
                    l.U.data = U_hat
//  [python]
def K_postprocess_step(self):

        for l in self.NN.layer:

            if hasattr(l,'lr') and l.lr:

                if not l.fixed:
                    
                    U_hat = torch.hstack((l.K[:,:l.dynamic_rank],l.U[:,:l.dynamic_rank]))

                    try:
                        U_hat,_ = torch.linalg.qr(U_hat)
                    except:
                        U_hat,_ = np.linalg.qr(U_hat)
                        U_hat = torch.tensor(U_hat)
                    l.U_hat[:,:2*l.dynamic_rank] = U_hat
                    l.M_hat[:2*l.dynamic_rank,:l.dynamic_rank] = l.U_hat[:,:2*l.dynamic_rank].T@l.U[:,:l.dynamic_rank]
                
                else:

                    try:
                        U_hat,_ = torch.linalg.qr(l.K)

                    except:
                        U_hat,_ = np.linalg.qr(U_hat)
                        U_hat = torch.tensor(U_hat)
                    l.M_hat.data = U_hat.T@l.U.data
                    l.U.data = U_hat
//  [python]
def _make_orthonorm_weights(self, Y: torch.Tensor) -> torch.Tensor:
        """
        Orthonormalize the output of the network using the Cholesky decomposition.

        Parameters
        ----------
        Y : torch.Tensor
            The output of the network.

        Returns
        -------
        torch.Tensor
            The orthonormalized output.

        Notes
        -----
        This function applies QR decomposition to orthonormalize the output (`Y`) of the network.
        The inverse of the R matrix is returned as the orthonormalization weights.
        """

        m = Y.shape[0]
        Q, R = torch.qr(Y)
        orthonorm_weights = np.sqrt(m) * torch.inverse(R)
        return orthonorm_weights
//  [python]
def _make_orthonorm_weights(self, Y: torch.Tensor) -> torch.Tensor:
        """
        Orthonormalize the output of the network using the Cholesky decomposition.

        Parameters
        ----------
        Y : torch.Tensor
            The output of the network.

        Returns
        -------
        torch.Tensor
            The orthonormalized output.

        Notes
        -----
        This function applies QR decomposition to orthonormalize the output (`Y`) of the network.
        The inverse of the R matrix is returned as the orthonormalization weights.
        """

        m = Y.shape[0]
        _, R = torch.linalg.qr(Y)
        orthonorm_weights = np.sqrt(m) * torch.inverse(R)
        return orthonorm_weights
//  [python]
def L_postprocess_step(self):

        for l in self.NN.layer:

            if hasattr(l,'lr') and l.lr:

                if not l.fixed:

                    V_hat = torch.hstack((l.L[:,:l.dynamic_rank],l.V[:,:l.dynamic_rank]))
                    try :
                        Q,R = torch.qr(V_hat)
                        V_hat = Q
                    except:
                        V_hat,_ = np.linalg.qr(V_hat.detach().numpy())
                        V_hat= torch.tensor(V_hat)
                    l.V_hat[:,:2*l.dynamic_rank] = V_hat
                    l.N_hat[:2*l.dynamic_rank,:l.dynamic_rank] = l.V_hat[:,:2*l.dynamic_rank].T@l.V[:,:l.dynamic_rank]

                else:

                    try :
                        Q,R = torch.qr(l.L)
                        V_hat = Q
                    except:
                        V_hat,_ = np.linalg.qr(V_hat.detach().numpy())
                        V_hat= torch.tensor(V_hat)
                    l.N_hat.data = V_hat.T@l.V.data
                    l.V.data = V_hat
//  [python]
def L_postprocess_step(self):

        for l in self.NN.layer:

            if hasattr(l,'lr') and l.lr:

                if not l.fixed:

                    V_hat = torch.hstack((l.L[:,:l.dynamic_rank],l.V[:,:l.dynamic_rank]))
                    try :
                        V_hat,_ = torch.linalg.qr(V_hat)
                    except:
                        V_hat,_ = np.linalg.qr(V_hat.detach().numpy())
                        V_hat= torch.tensor(V_hat)
                    l.V_hat[:,:2*l.dynamic_rank] = V_hat
                    l.N_hat[:2*l.dynamic_rank,:l.dynamic_rank] = l.V_hat[:,:2*l.dynamic_rank].T@l.V[:,:l.dynamic_rank]

                else:

                    try :
                        V_hat,_ = torch.linalg.qr(l.L)
                    except:
                        V_hat,_ = np.linalg.qr(V_hat.detach().numpy())
                        V_hat= torch.tensor(V_hat)
                    l.N_hat.data = V_hat.T@l.V.data
                    l.V.data = V_hat
//  [python]
def solve(self, x):
        batch, m = x.shape[:2]
        if (self.num_eigen_values <= 0) or (self.num_eigen_values > m): self.num_eigen_values = m
        Q, R = torch.qr(x)
        if self.enable_stop_condition: Q_his = Q[:, :, :self.num_eigen_values].detach()

        for _ in range(self.num_iters):
            x_t = torch.einsum('bmn,bnk->bmk', x, Q)
            Q, R = torch.qr(x_t)

            if self.enable_stop_condition:
                cur_Q = Q[:, :, :self.num_eigen_values].detach()
                if self.stop_condition(cur_Q, Q_his): break
                Q_his = cur_Q

        l, sort_indices = torch.sort(torch.diagonal(R, dim1=1, dim2=2), dim=1, descending=True)
        l = l[:, :self.num_eigen_values]
        Q = Q[
            torch.arange(batch).view(batch, 1, 1), torch.arange(m).view(1, m, 1),
            sort_indices.view(batch, 1, m)[:, :, :self.num_eigen_values]]
        Q = self.uniform_solution_direction(Q)

        return l, Q
//  [python]
def solve(self, x):
        batch, m = x.shape[:2]
        if (self.num_eigen_values <= 0) or (self.num_eigen_values > m): self.num_eigen_values = m
        Q, R = torch.linalg.qr(x)
        if self.enable_stop_condition: Q_his = Q[:, :, :self.num_eigen_values].detach()

        for _ in range(self.num_iters):
            x_t = torch.einsum('bmn,bnk->bmk', x, Q)
            Q, R = torch.linalg.qr(x_t)

            if self.enable_stop_condition:
                cur_Q = Q[:, :, :self.num_eigen_values].detach()
                if self.stop_condition(cur_Q, Q_his): break
                Q_his = cur_Q

        l, sort_indices = torch.sort(torch.diagonal(R, dim1=1, dim2=2), dim=1, descending=True)
        l = l[:, :self.num_eigen_values]
        Q = Q[
            torch.arange(batch).view(batch, 1, 1), torch.arange(m).view(1, m, 1),
            sort_indices.view(batch, 1, m)[:, :, :self.num_eigen_values]]
        Q = self.uniform_solution_direction(Q)

        return l, Q
//  [python]
def qr(data, meta, sizes):
    Qdata = torch.zeros((sizes[0],), dtype=data.dtype, device=data.device)
    Rdata = torch.zeros((sizes[1],), dtype=data.dtype, device=data.device)
    for (sl, D, slQ, DQ, slR, DR) in meta:
        Q, R = torch.qr(data[slice(*sl)].view(D))
        sR = torch.sign(real(R.diag()))
        sR[sR == 0] = 1
        Qdata[slice(*slQ)].view(DQ)[:] = Q * sR  # positive diag of R
        Rdata[slice(*slR)].view(DR)[:] = sR.reshape([-1, 1]) * R
    return Qdata, Rdata
//  [python]
def qr(data, meta, sizes):
    Qdata = torch.zeros((sizes[0],), dtype=data.dtype, device=data.device)
    Rdata = torch.zeros((sizes[1],), dtype=data.dtype, device=data.device)
    for (sl, D, slQ, DQ, slR, DR) in meta:
        Q, R = torch.linalg.qr(data[slice(*sl)].view(D))
        sR = torch.sign(real(R.diag()))
        sR[sR == 0] = 1
        Qdata[slice(*slQ)].view(DQ)[:] = Q * sR  # positive diag of R
        Rdata[slice(*slR)].view(DR)[:] = sR.reshape([-1, 1]) * R
    return Qdata, Rdata
//  [python]
def generate_random_state(self):
        self.random_index = torch.tensor([i for i in torch.randperm(self.input_dim)])
        self.random_subsets = torch.chunk(self.random_index, self.subset_num, dim=0)
        self.selection_feature = []
        for subset_i in range(len(self.random_subsets)):
            vector_index = nn.Parameter(
                torch.zeros(self.input_dim), requires_grad=False
            )
            non_zero_f = sorted(self.random_subsets[subset_i])
            for j in range(len(non_zero_f)):
                channel = non_zero_f[j]
                vector_index[channel] = 1
            self.selection_feature.append(vector_index)
        self.selection_feature = nn.Parameter(
            torch.stack(self.selection_feature, dim=0), requires_grad=False
        )
        self.shuffle_index1 = nn.Parameter(
            torch.randperm(self.input_dim), requires_grad=False
        )
        self.shuffle_index2 = nn.Parameter(
            torch.randperm(self.input_dim), requires_grad=False
        )
        self.matrix = torch.randn(self.input_dim, self.input_dim)
        self.q, self.r = torch.qr(self.matrix)
        self.q[:, -1] = 0
        self.q = nn.Parameter(self.q, requires_grad=False)
//  [python]
def generate_random_state(self):
        self.random_index = torch.tensor([i for i in torch.randperm(self.input_dim)])
        self.random_subsets = torch.chunk(self.random_index, self.subset_num, dim=0)
        self.selection_feature = []
        for subset_i in range(len(self.random_subsets)):
            vector_index = nn.Parameter(
                torch.zeros(self.input_dim), requires_grad=False
            )
            non_zero_f = sorted(self.random_subsets[subset_i])
            for j in range(len(non_zero_f)):
                channel = non_zero_f[j]
                vector_index[channel] = 1
            self.selection_feature.append(vector_index)
        self.selection_feature = nn.Parameter(
            torch.stack(self.selection_feature, dim=0), requires_grad=False
        )
        self.shuffle_index1 = nn.Parameter(
            torch.randperm(self.input_dim), requires_grad=False
        )
        self.shuffle_index2 = nn.Parameter(
            torch.randperm(self.input_dim), requires_grad=False
        )
        self.matrix = torch.randn(self.input_dim, self.input_dim)
        self.q, self.r = torch.linalg.qr(self.matrix)
        self.q[:, -1] = 0
        self.q = nn.Parameter(self.q, requires_grad=False)
//  [python]
def random_ortho(n, m, std):
    q, _ = torch.qr(torch.randn(n, m) * std)
    return q
//  [python]
def random_ortho(n, m, std):
    q, _ = torch.linalg.qr(torch.randn(n, m) * std)
    return q
//  [python]
bad = torch.qr()
good = torch.qr()
//  [python]
bad = torch.qr()
good = torch.linalg.qr()
//  [python]
user_weights = np.reshape(np.ones(nr_of_users*nr_of_samples_per_batch),(nr_of_samples_per_batch,nr_of_users,1))
user_weights = torch.as_tensor(user_weights)
user_weights_for_regular_WMMSE = np.ones(nr_of_users)
#Q = torch.randn(32,32,dtype=torch.float)
Q = torch.rand(32,32,dtype=torch.float)
q, _ = torch.qr(Q)
//  [python]
user_weights = np.reshape(np.ones(nr_of_users*nr_of_samples_per_batch),(nr_of_samples_per_batch,nr_of_users,1))
user_weights = torch.as_tensor(user_weights)
user_weights_for_regular_WMMSE = np.ones(nr_of_users)
#Q = torch.randn(32,32,dtype=torch.float)
Q = torch.rand(32,32,dtype=torch.float)
q, _ = torch.linalg.qr(Q)
//  [python]
def _decay_helper(
        svals,
        shape=(100, 100),
        rank=10,
        decay=0.5,
        symmetric=True,
        seed=0b1110101001010101011,
        dtype=torch.float64,
        device="cpu",
    ):
        """Given singular values, produces a random matrix.

        Helper method for the polynomial and exp-decay matrices. Given the
        singular values, it samples the left- and right- singular spaces,
        and returns the final matrix.
        """
        min_shape = min(shape)
        # build singular bases using QR subgroup algorithm (Diaconis). QR is not
        # fastest, but these are test matrices so speed is not crucial.
        G_left = gaussian_noise(
            (shape[0], min_shape),
            mean=0.0,
            std=1.0,
            seed=seed,
            dtype=dtype,
            device=device,
        )
        U, _ = torch.qr(G_left)
        del G_left
        #
        if symmetric:
            result = U @ torch.diag(svals) @ U.T
        else:
            G_right = gaussian_noise(
                (shape[1], min_shape),
                mean=0.0,
                std=1.0,
                seed=seed + 1,
                dtype=dtype,
                device=device,
            )
            V, _ = torch.qr(G_right)
            result = U @ torch.diag(svals) @ V.T
        #
        return result
//  [python]
def _decay_helper(
        svals,
        shape=(100, 100),
        rank=10,
        decay=0.5,
        symmetric=True,
        seed=0b1110101001010101011,
        dtype=torch.float64,
        device="cpu",
    ):
        """Given singular values, produces a random matrix.

        Helper method for the polynomial and exp-decay matrices. Given the
        singular values, it samples the left- and right- singular spaces,
        and returns the final matrix.
        """
        min_shape = min(shape)
        # build singular bases using QR subgroup algorithm (Diaconis). QR is not
        # fastest, but these are test matrices so speed is not crucial.
        G_left = gaussian_noise(
            (shape[0], min_shape),
            mean=0.0,
            std=1.0,
            seed=seed,
            dtype=dtype,
            device=device,
        )
        U, _ = torch.linalg.qr(G_left)
        del G_left
        #
        if symmetric:
            result = U @ torch.diag(svals) @ U.T
        else:
            G_right = gaussian_noise(
                (shape[1], min_shape),
                mean=0.0,
                std=1.0,
                seed=seed + 1,
                dtype=dtype,
                device=device,
            )
            V, _ = torch.linalg.qr(G_right)
            result = U @ torch.diag(svals) @ V.T
        #
        return result
//  [python]
def qr(batched_matrix, positive_diag=True) -> Tuple[Tensor, Tensor]:
    """QR decomposition with the option of making the diagonal elements
    of R positive.

    Args:
        batched_matrix (Tensor): An N-D tensor with N > 2.
        positive_diag (bool, optional): Whether to make the diagonal elements
            of R positive. Defaults to True.

    Returns:
        Tuple[Tensor, Tensor]: The Q and R.
    """
    Q, R = torch.qr(batched_matrix)
    if positive_diag:
        # The QR decomposition returned by pytorch does not guarantee
        # that the matrix R has positive diagonal elements.
        # To ensure that, we first construct a signature matrix S
        # which is a diagonal matrix with diagonal elements +1 or -1
        # depending on the sign of the corresponding diagonal element in R.
        # If we premultiply R with S, we ensure that the diagonal is positive;
        # however, the corresponding Q needs to be changed now to ensure that
        # that we get back the original matrix when we multiply the new Q and
        # the new R. This can be done by postmultiplying Q with S because
        # then we get (Q @ S) @ (S @ R) = Q @ R as S @ S = I, where (Q @ S)
        # and (S @ R) are the new Q and new R respectively.
        R_diag_sign = torch.sgn(torch.diagonal(R, dim1=-2, dim2=-1))
        S = torch.diag_embed(R_diag_sign, dim1=-2, dim2=-1)
        Q = Q @ S
        R = S @ R
    return Q, R
//  [python]
def qr(batched_matrix, positive_diag=True) -> Tuple[Tensor, Tensor]:
    """QR decomposition with the option of making the diagonal elements
    of R positive.

    Args:
        batched_matrix (Tensor): An N-D tensor with N > 2.
        positive_diag (bool, optional): Whether to make the diagonal elements
            of R positive. Defaults to True.

    Returns:
        Tuple[Tensor, Tensor]: The Q and R.
    """
    Q, R = torch.linalg.qr(batched_matrix)
    if positive_diag:
        # The QR decomposition returned by pytorch does not guarantee
        # that the matrix R has positive diagonal elements.
        # To ensure that, we first construct a signature matrix S
        # which is a diagonal matrix with diagonal elements +1 or -1
        # depending on the sign of the corresponding diagonal element in R.
        # If we premultiply R with S, we ensure that the diagonal is positive;
        # however, the corresponding Q needs to be changed now to ensure that
        # that we get back the original matrix when we multiply the new Q and
        # the new R. This can be done by postmultiplying Q with S because
        # then we get (Q @ S) @ (S @ R) = Q @ R as S @ S = I, where (Q @ S)
        # and (S @ R) are the new Q and new R respectively.
        R_diag_sign = torch.sgn(torch.diagonal(R, dim1=-2, dim2=-1))
        S = torch.diag_embed(R_diag_sign, dim1=-2, dim2=-1)
        Q = Q @ S
        R = S @ R
    return Q, R
//  [python]
def generate_observations(
    sample_size,
    mean_vector: torch.tensor,
    uv_scale: torch.tensor,
) -> torch.tensor:
    """
    Generate independent random observations
    """
    if mean_vector.shape != uv_scale.shape:
        raise ValueError("Vectors must conform")

    torch.manual_seed(MANUAL_SEED)
    white_noise = torch.randn((sample_size, mean_vector.shape[0]))
    # Correct sample mean, sample variance, and orthogonality slightly so that the
    # sample statistics are what we want.
    white_noise = white_noise - torch.mean(white_noise, dim=0)
    q, r = torch.qr(white_noise)
    white_noise = q * sqrt(sample_size)

    # Resulting sample should have 1) zero mean; 2) unit variance; 3) orthogonal columns

    # Multiply and drop the last dimension.
    random_observations = uv_scale * white_noise + mean_vector

    return random_observations, white_noise
//  [python]
def generate_observations(
    sample_size,
    mean_vector: torch.tensor,
    uv_scale: torch.tensor,
) -> torch.tensor:
    """
    Generate independent random observations
    """
    if mean_vector.shape != uv_scale.shape:
        raise ValueError("Vectors must conform")

    torch.manual_seed(MANUAL_SEED)
    white_noise = torch.randn((sample_size, mean_vector.shape[0]))
    # Correct sample mean, sample variance, and orthogonality slightly so that the
    # sample statistics are what we want.
    white_noise = white_noise - torch.mean(white_noise, dim=0)
    white_noise = torch.linalg.qr(white_noise, mode="reduced")[0] * sqrt(sample_size)

    # Resulting sample should have 1) zero mean; 2) unit variance; 3) orthogonal columns

    # Multiply and drop the last dimension.
    random_observations = uv_scale * white_noise + mean_vector

    return random_observations, white_noise
//  [python]
def __init__(self, args):
        self.args = args
        self.device = torch.device(0)

        self.biased_classifier = SimpleConv64(
            latent_dim=1, num_channels=args.num_channels, image_size=args.image_size).to(self.device)
        self.biased_classifier.load_state_dict(
            torch.load(args.biased_classifier_ckpt)['model'])
        print('biased classifier checkpoint loaded.')
        self.biased_classifier.eval()

        model_cl = getattr(models, args.alg)
        self.model = model_cl(args)
        self.model.load_checkpoint(args.ckpt_load)
        self.model.net_mode(train=False)
        print('generative model checkpoint loaded.')

        w = torch.randn(1, args.z_dim, device=self.device)
        b = torch.randn(1, 1, device=self.device)

        self.w = torch.nn.Parameter(w, requires_grad=True)
        self.b = torch.nn.Parameter(b, requires_grad=True)
        if args.optimizer == 'SGD':
            self.optimizer = torch.optim.SGD([self.w, self.b], args.lr)
        elif args.optimizer == 'Adam':
            self.optimizer = torch.optim.Adam(
                [self.w, self.b], args.lr, betas=(args.beta1, args.beta2))
        else:
            raise NotImplementedError

        hyperplanes = np.load(args.normal_vector_npz_fpath)
        gt_normal_vector_mat = torch.from_numpy(
            hyperplanes['normal_vector_mat']).float()
        q, r = torch.qr(gt_normal_vector_mat)
        gt_orth_normal_vector_mat = q
        assert gt_orth_normal_vector_mat.shape == gt_normal_vector_mat.shape
        self.gt_biased_normal_vec = gt_orth_normal_vector_mat[:, args.bias_attr_index].to(
            self.device).reshape(1, -1)
        self.gt_target_normal_vec = gt_orth_normal_vector_mat[:, args.target_attr_index].to(
            self.device).reshape(1, -1)
        gt_offset_mat = torch.from_numpy(hyperplanes['offset_mat'])
        self.gt_biased_offset = gt_offset_mat[args.bias_attr_index].to(
            self.device)

        gt_other_normal_vectors = gt_normal_vector_mat[:, [i for i in range(
            gt_normal_vector_mat.shape[1]) if i != args.bias_attr_index]].to(self.device)
        q, r = torch.qr(gt_other_normal_vectors)
        gt_orth_other_normal_vectors = q
        assert gt_orth_other_normal_vectors.shape == gt_other_normal_vectors.shape

        self.unit_gt_other_normal_vectors = gt_orth_other_normal_vectors / \
            torch.norm(gt_orth_other_normal_vectors, dim=0).unsqueeze(0)

        self.traverse_steps = torch.arange(
            args.traverse_min, args.traverse_max, args.traverse_spacing, device=self.device).reshape(1, -1, 1)
//  [python]
def __init__(self, args):
        self.args = args
        self.device = torch.device(0)

        self.biased_classifier = SimpleConv64(
            latent_dim=1, num_channels=args.num_channels, image_size=args.image_size).to(self.device)
        self.biased_classifier.load_state_dict(
            torch.load(args.biased_classifier_ckpt)['model'])
        print('biased classifier checkpoint loaded.')
        self.biased_classifier.eval()

        model_cl = getattr(models, args.alg)
        self.model = model_cl(args)
        self.model.load_checkpoint(args.ckpt_load)
        self.model.net_mode(train=False)
        print('generative model checkpoint loaded.')

        w = torch.randn(1, args.z_dim, device=self.device)
        b = torch.randn(1, 1, device=self.device)

        self.w = torch.nn.Parameter(w, requires_grad=True)
        self.b = torch.nn.Parameter(b, requires_grad=True)
        if args.optimizer == 'SGD':
            self.optimizer = torch.optim.SGD([self.w, self.b], args.lr)
        elif args.optimizer == 'Adam':
            self.optimizer = torch.optim.Adam(
                [self.w, self.b], args.lr, betas=(args.beta1, args.beta2))
        else:
            raise NotImplementedError

        hyperplanes = np.load(args.normal_vector_npz_fpath)
        gt_normal_vector_mat = torch.from_numpy(
            hyperplanes['normal_vector_mat']).float()
        gt_orth_normal_vector_mat = torch.linalg.qr(gt_normal_vector_mat)[0]
        assert gt_orth_normal_vector_mat.shape == gt_normal_vector_mat.shape
        self.gt_biased_normal_vec = gt_orth_normal_vector_mat[:, args.bias_attr_index].to(
            self.device).reshape(1, -1)
        self.gt_target_normal_vec = gt_orth_normal_vector_mat[:, args.target_attr_index].to(
            self.device).reshape(1, -1)
        gt_offset_mat = torch.from_numpy(hyperplanes['offset_mat'])
        self.gt_biased_offset = gt_offset_mat[args.bias_attr_index].to(
            self.device)

        gt_other_normal_vectors = gt_normal_vector_mat[:, [i for i in range(
            gt_normal_vector_mat.shape[1]) if i != args.bias_attr_index]].to(self.device)
        gt_orth_other_normal_vectors = torch.qr(gt_other_normal_vectors)[0]
        assert gt_orth_other_normal_vectors.shape == gt_other_normal_vectors.shape

        self.unit_gt_other_normal_vectors = gt_orth_other_normal_vectors / \
            torch.norm(gt_orth_other_normal_vectors, dim=0).unsqueeze(0)

        self.traverse_steps = torch.arange(
            args.traverse_min, args.traverse_max, args.traverse_spacing, device=self.device).reshape(1, -1, 1)
//  [python]
def random_unitary(n, device: Optional[torch.device] = None):
    """
    Returns random unitary matrix from the Haar measure of size n x n.
    
    Unitary matrix is created as described in this `paper
    <https://arxiv.org/abs/math-ph/0609050v2>`_.
    """
    mat = torch.randn(n, n, device=device)
    q, r = torch.qr(mat)
    d = torch.diagonal(r)
    ph = d / d.abs()
    q = q @ torch.diag(ph)
    return q
//  [python]
def random_unitary(n, device: Optional[torch.device] = None):
    """
    Returns random unitary matrix from the Haar measure of size n x n.
    
    Unitary matrix is created as described in this `paper
    <https://arxiv.org/abs/math-ph/0609050v2>`_.
    """
    mat = torch.randn(n, n, device=device)
    q, r = torch.linalg.qr(mat)
    d = torch.diagonal(r)
    ph = d / d.abs()
    q = q @ torch.diag(ph)
    return q
//  [python]
def init_weights(self):
        with torch.no_grad():
            w = self.fc.weight
            nn.init.normal_(w)
            wt = w.reshape(-1, self.subdim,
                           self.in_channels).permute([0, 2, 1])
            wt = torch.qr(wt)[0]
            w = wt.permute([0, 2, 1]).reshape(-1, self.in_channels)
            self.fc.weight.add_(w - self.fc.weight)
//  [python]
def init_weights(self):
        with torch.no_grad():
            w = self.fc.weight
            nn.init.normal_(w)
            wt = w.reshape(-1, self.subdim,
                           self.in_channels).permute([0, 2, 1])
            wt = torch.linalg.qr(wt).Q
            w = wt.permute([0, 2, 1]).reshape(-1, self.in_channels)
            self.fc.weight.add_(w - self.fc.weight)
//  [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)
    # Originally used the deprecated torch.qr(unstructured_block.cpu())
    q, r = torch.linalg.qr(unstructured_block.cpu())
    q, r = map(lambda t: t.to(device), (q, r))
    return q.t()