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

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

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

    rows = tensor.size(0)
    cols = tensor.numel() // rows
    flattened = tensor.new(rows, cols).normal_(0, 1)

    if rows < cols:
        flattened.t_()

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

    if rows < cols:
        q.t_()

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

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

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

    rows = tensor.size(0)
    cols = tensor.numel() // rows
    flattened = tensor.new(rows, cols).normal_(0, 1)

    if rows < cols:
        flattened.t_()

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

    if rows < cols:
        q.t_()

    with torch.no_grad():
        tensor.view_as(q).copy_(q)
        tensor.mul_(gain)
    return tensor
//  [python]
def __init__(self, c, cache_inverse=False):
        super(Invertible1x1Conv, self).__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.cache_inverse = cache_inverse
//  [python]
def __init__(self, c, cache_inverse=False):
        super(Invertible1x1Conv, self).__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.cache_inverse = cache_inverse
//  [python]
def sample_Q(self, num_channels, n_sample):
        Qs = []
        errs = []
        for i in range(n_sample):
            A = torch.rand(num_channels, num_channels, dtype=torch.float32)
            Q, R = torch.qr(A)
            Qs.append(Q)
            err = torch.max(torch.abs(torch.inverse(torch.inverse(Q)) - Q)).item()
            errs.append(err)
        idx_best = np.argmin(errs)
        Q = Qs[idx_best]
        return Q, torch.inverse(Q)
//  [python]
def sample_Q(self, num_channels, n_sample):
        Qs = []
        errs = []
        for i in range(n_sample):
            A = torch.rand(num_channels, num_channels, dtype=torch.float32)
            Q, R = torch.linalg.qr(A)
            Qs.append(Q)
            err = torch.max(torch.abs(torch.inverse(torch.inverse(Q)) - Q)).item()
            errs.append(err)
        idx_best = np.argmin(errs)
        Q = Qs[idx_best]
        return Q, torch.inverse(Q)
//  [python]
def orthogonal_random_matrix_(rows, columns, device):
    """Generate a random matrix whose columns are orthogonal to each
    other (in groups of size `rows`) and their norms is drawn from the
    chi-square distribution with `rows` degrees of freedom
    (namely the norm of a `rows`-dimensional vector distributed as N(0, I)).
    """
    w = torch.zeros([rows, columns], device=device)
    start = 0
    while start < columns:
        end = min(start+rows, columns)
        block = torch.randn(rows, rows, device=device)
        norms = torch.sqrt(torch.einsum("ab,ab->a", block, block))
        Q, _ = torch.qr(block)  # Q is orthonormal
        
        w[:, start:end] = (
            Q[:, :end-start] * norms[None, :end-start]
        )
        start += rows
    return w
//  [python]
def orthogonal_random_matrix_(rows, columns, device):
    """Generate a random matrix whose columns are orthogonal to each
    other (in groups of size `rows`) and their norms is drawn from the
    chi-square distribution with `rows` degrees of freedom
    (namely the norm of a `rows`-dimensional vector distributed as N(0, I)).
    """
    w = torch.zeros([rows, columns], device=device)
    start = 0
    while start < columns:
        end = min(start+rows, columns)
        block = torch.randn(rows, rows, device=device)
        norms = torch.sqrt(torch.einsum("ab,ab->a", block, block))
        Q, _ = torch.linalg.qr(block)  # Q is orthonormal
        
        w[:, start:end] = (
            Q[:, :end-start] * norms[None, :end-start]
        )
        start += rows
    return w
//  [python]
def __init__(self, channels, n_split=4, no_jacobian=False, **kwargs):
    super().__init__()
    assert(n_split % 2 == 0)
    self.channels = channels
    self.n_split = n_split
    self.no_jacobian = no_jacobian
    
    w_init = torch.qr(torch.FloatTensor(self.n_split, self.n_split).normal_())[0]
    if torch.det(w_init) < 0:
      w_init[:,0] = -1 * w_init[:,0]
    self.weight = nn.Parameter(w_init)
//  [python]
def __init__(self, channels, n_split=4, no_jacobian=False, **kwargs):
    super().__init__()
    assert(n_split % 2 == 0)
    self.channels = channels
    self.n_split = n_split
    self.no_jacobian = no_jacobian
    
    w_init = torch.linalg.qr(torch.FloatTensor(self.n_split, self.n_split).normal_())[0]
    if torch.det(w_init) < 0:
      w_init[:,0] = -1 * w_init[:,0]
    self.weight = nn.Parameter(w_init)
//  [python]
def cca(features_x, features_y):
    """Compute the mean squared CCA correlation (R^2_{CCA}).

    Args:
        features_x: A num_examples x num_features matrix of features.
        features_y: A num_examples x num_features matrix of features.

    Returns:
        The mean squared CCA correlations between X and Y.
    """
    qx, _ = torch.qr(features_x)  # Or use SVD with full_matrices=False.
    qy, _ = torch.qr(features_y)
    return torch.linalg.norm(qx.T @ (qy)) ** 2 / min(features_x.shape[1], features_y.shape[1])
//  [python]
def cca(features_x, features_y):
    """Compute the mean squared CCA correlation (R^2_{CCA}).

    Args:
        features_x: A num_examples x num_features matrix of features.
        features_y: A num_examples x num_features matrix of features.

    Returns:
        The mean squared CCA correlations between X and Y.
    """
    qx, _ = torch.linalg.qr(features_x)  # Or use SVD with full_matrices=False.
    qy, _ = torch.linalg.qr(features_y)
    return torch.linalg.norm(qx.T @ (qy)) ** 2 / min(features_x.shape[1], features_y.shape[1])
//  [python]
def __init__(self, n_channels=3, lu_factorize=False):
        super().__init__()
        self.lu_factorize = lu_factorize

        # initiaize a 1x1 convolution weight matrix
        w = torch.randn(n_channels, n_channels)
        w = torch.qr(w)[0]  # note: nn.init.orthogonal_ returns orth matrices with dets +/- 1 which complicates the inverse call below

        if lu_factorize:
            # compute LU factorization
            p, l, u = torch.btriunpack(*w.unsqueeze(0).btrifact())
            # initialize model parameters
            self.p, self.l, self.u = nn.Parameter(p.squeeze()), nn.Parameter(l.squeeze()), nn.Parameter(u.squeeze())
            s = self.u.diag()
            self.log_s = nn.Parameter(s.abs().log())
            self.register_buffer('sign_s', s.sign())  # note: not optimizing the sign; det W remains the same sign
            self.register_buffer('l_mask', torch.tril(torch.ones_like(self.l), -1))  # store mask to compute LU in forward/inverse pass
        else:
            self.w = nn.Parameter(w)
//  [python]
def __init__(self, n_channels=3, lu_factorize=False):
        super().__init__()
        self.lu_factorize = lu_factorize

        # initiaize a 1x1 convolution weight matrix
        w = torch.randn(n_channels, n_channels)
        w = torch.linalg.qr(w)[0]  # note: nn.init.orthogonal_ returns orth matrices with dets +/- 1 which complicates the inverse call below

        if lu_factorize:
            # compute LU factorization
            p, l, u = torch.btriunpack(*w.unsqueeze(0).btrifact())
            # initialize model parameters
            self.p, self.l, self.u = nn.Parameter(p.squeeze()), nn.Parameter(l.squeeze()), nn.Parameter(u.squeeze())
            s = self.u.diag()
            self.log_s = nn.Parameter(s.abs().log())
            self.register_buffer('sign_s', s.sign())  # note: not optimizing the sign; det W remains the same sign
            self.register_buffer('l_mask', torch.tril(torch.ones_like(self.l), -1))  # store mask to compute LU in forward/inverse pass
        else:
            self.w = nn.Parameter(w)
//  [python]
def __init__(self, num_channels):
        super().__init__()
        w_shape = [num_channels, num_channels]
        # initialize with random orhogonal matrix whose log det is 0
        w_init = torch.qr(torch.randn(*w_shape))[0]
        self.weight = nn.Parameter(torch.Tensor(w_init))
        self.w_shape = w_shape
//  [python]
def __init__(self, num_channels):
        super().__init__()
        w_shape = [num_channels, num_channels]
        # initialize with random orhogonal matrix whose log det is 0
        w_init = torch.linalg.qr(torch.randn(*w_shape))[0]
        self.weight = nn.Parameter(torch.Tensor(w_init))
        self.w_shape = w_shape
//  [python]
def __init__(self, dim, ones=True, scramble=False, hetero=True, hidden=False):
        self.hetero = hetero
        self.hidden = hidden
        self.dim = dim // 2

        if ones:
            self.wxy = torch.eye(self.dim)
            self.wyz = torch.eye(self.dim)
        else:
            self.wxy = torch.randn(self.dim, self.dim) / dim
            self.wyz = torch.randn(self.dim, self.dim) / dim

        if scramble:
            self.scramble, _ = torch.qr(torch.randn(dim, dim))
        else:
            self.scramble = torch.eye(dim)

        if hidden:
            self.whx = torch.randn(self.dim, self.dim) / dim
            self.why = torch.randn(self.dim, self.dim) / dim
            self.whz = torch.randn(self.dim, self.dim) / dim
        else:
            self.whx = torch.eye(self.dim, self.dim)
            self.why = torch.zeros(self.dim, self.dim)
            self.whz = torch.zeros(self.dim, self.dim)
//  [python]
def __init__(self, dim, ones=True, scramble=False, hetero=True, hidden=False):
        self.hetero = hetero
        self.hidden = hidden
        self.dim = dim // 2

        if ones:
            self.wxy = torch.eye(self.dim)
            self.wyz = torch.eye(self.dim)
        else:
            self.wxy = torch.randn(self.dim, self.dim) / dim
            self.wyz = torch.randn(self.dim, self.dim) / dim

        if scramble:
            self.scramble, _ = torch.linalg.qr(torch.randn(dim, dim))
        else:
            self.scramble = torch.eye(dim)

        if hidden:
            self.whx = torch.randn(self.dim, self.dim) / dim
            self.why = torch.randn(self.dim, self.dim) / dim
            self.whz = torch.randn(self.dim, self.dim) / dim
        else:
            self.whx = torch.eye(self.dim, self.dim)
            self.why = torch.zeros(self.dim, self.dim)
            self.whz = torch.zeros(self.dim, self.dim)
//  [python]
def __init__(self, in_channel):
        super().__init__()

        weight = torch.randn(in_channel, in_channel).to(device)
        q, _ = torch.qr(weight)  # A = QR with Q being an orthogonal matrix and R being an upper triangular matrix
        weight = q.unsqueeze(2).unsqueeze(3)  # add dim 2,3 into (C,C,1,1)
        self.weight = nn.Parameter(weight)
//  [python]
def __init__(self, in_channel):
        super().__init__()

        weight = torch.randn(in_channel, in_channel).to(device)
        q, _ = torch.linalg.qr(weight)  # A = QR with Q being an orthogonal matrix and R being an upper triangular matrix
        weight = q.unsqueeze(2).unsqueeze(3)  # add dim 2,3 into (C,C,1,1)
        self.weight = nn.Parameter(weight)
//  [python]
def basis(A):
    """Return orthogonal basis of A columns.
    """
    if A.is_cuda:
        # torch.orgqr is not available in CUDA
        Q, _ = torch.qr(A, some=True)
    else:
        Q = torch.orgqr(*torch.geqrf(A))
    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, mode='reduced')
    else:
        Q = torch.orgqr(*torch.linalg.qr(A, mode='raw'))
    return Q
//  [python]
def __init__(self, dim, scramble=False, hetero=True, hidden=False):
        self.hetero = hetero
        self.hidden = hidden
        self.dim = dim // 2
        ones = True

        if ones:
            self.wxy = torch.eye(self.dim)
            self.wyz = torch.eye(self.dim)
        else:
            self.wxy = torch.randn(self.dim, self.dim) / dim
            self.wyz = torch.randn(self.dim, self.dim) / dim

        if scramble:
            self.scramble, _ = torch.qr(torch.randn(dim, dim))
        else:
            self.scramble = torch.eye(dim)

        if hidden:
            self.whx = torch.randn(self.dim, self.dim) / dim
            self.why = torch.randn(self.dim, self.dim) / dim
            self.whz = torch.randn(self.dim, self.dim) / dim
        else:
            self.whx = torch.eye(self.dim, self.dim)
            self.why = torch.zeros(self.dim, self.dim)
            self.whz = torch.zeros(self.dim, self.dim)
//  [python]
def __init__(self, dim, scramble=False, hetero=True, hidden=False):
        self.hetero = hetero
        self.hidden = hidden
        self.dim = dim // 2
        ones = True

        if ones:
            self.wxy = torch.eye(self.dim)
            self.wyz = torch.eye(self.dim)
        else:
            self.wxy = torch.randn(self.dim, self.dim) / dim
            self.wyz = torch.randn(self.dim, self.dim) / dim

        if scramble:
            self.scramble, _ = torch.linalg.qr(torch.randn(dim, dim))
        else:
            self.scramble = torch.eye(dim)

        if hidden:
            self.whx = torch.randn(self.dim, self.dim) / dim
            self.why = torch.randn(self.dim, self.dim) / dim
            self.whz = torch.randn(self.dim, self.dim) / dim
        else:
            self.whx = torch.eye(self.dim, self.dim)
            self.why = torch.zeros(self.dim, self.dim)
            self.whz = torch.zeros(self.dim, self.dim)
//  [python]
def __init__(self, dim, args, scramble=False, hetero=True, hidden=False):
        self.hetero = hetero
        self.hidden = hidden
        self.dim = dim // 2 # UGLY!
        ones = args["env_ones"]
        self.args = args

        if ones:
            self.wxy = torch.eye(self.dim)
            self.wyz = torch.eye(self.dim)
        else:
            self.wxy = torch.randn(self.dim, self.dim) / dim
            self.wyz = torch.randn(self.dim, self.dim) / dim

        if scramble:
            self.scramble, _ = torch.qr(torch.randn(dim, dim))
        else:
            self.scramble = torch.eye(dim)

        if hidden:
            self.whx = torch.randn(self.dim, self.dim) / dim
            self.why = torch.randn(self.dim, self.dim) / dim
            self.whz = torch.randn(self.dim, self.dim) / dim
        else:
            self.whx = torch.eye(self.dim, self.dim)
            self.why = torch.zeros(self.dim, self.dim)
            self.whz = torch.zeros(self.dim, self.dim)
//  [python]
def __init__(self, dim, args, scramble=False, hetero=True, hidden=False):
        self.hetero = hetero
        self.hidden = hidden
        self.dim = dim // 2 # UGLY!
        ones = args["env_ones"]
        self.args = args

        if ones:
            self.wxy = torch.eye(self.dim)
            self.wyz = torch.eye(self.dim)
        else:
            self.wxy = torch.randn(self.dim, self.dim) / dim
            self.wyz = torch.randn(self.dim, self.dim) / dim

        if scramble:
            self.scramble, _ = torch.linalg.qr(torch.randn(dim, dim))
        else:
            self.scramble = torch.eye(dim)

        if hidden:
            self.whx = torch.randn(self.dim, self.dim) / dim
            self.why = torch.randn(self.dim, self.dim) / dim
            self.whz = torch.randn(self.dim, self.dim) / dim
        else:
            self.whx = torch.eye(self.dim, self.dim)
            self.why = torch.zeros(self.dim, self.dim)
            self.whz = torch.zeros(self.dim, self.dim)
//  [python]
def orthogonalize_tensor(tensor):
    assert len(tensor.shape) == 2
    # flattened = tensor.new(rows, cols).normal_(0, 1)

    # Compute the qr factorization
    q, r = torch.qr(tensor)
    # Make Q uniform according to https://arxiv.org/pdf/math-ph/0609050.pdf
    d = torch.diag(r, 0)
    ph = d.sign()
    q *= ph
    tensor.view_as(q).copy_(q)
    return tensor
//  [python]
def orthogonalize_tensor(tensor):
    assert len(tensor.shape) == 2
    # flattened = tensor.new(rows, cols).normal_(0, 1)

    # Compute the qr factorization
    q, r = torch.linalg.qr(tensor)
    # Make Q uniform according to https://arxiv.org/pdf/math-ph/0609050.pdf
    d = torch.diag(r, 0)
    ph = d.sign()
    q *= ph
    tensor.view_as(q).copy_(q)
    return tensor
//  [python]
def __init__(self, ndim, K=None, M=200):

        #K: number of slices per iteration. The same K in max K-SWD.
        #M: number of spline knots of rational quadratic spline

        super().__init__()
        self.ndim = ndim
        if K is None:
            self.K = ndim
        else:
            self.K = K
        self.M = M

        ATi = torch.randn(self.ndim, self.K)
        Q, R = torch.qr(ATi)
        L = torch.sign(torch.diag(R))
        A = (Q * L)

        self.A = nn.Parameter(A)
        self.transform1D = RQspline(self.K, M)
//  [python]
def __init__(self, ndim, K=None, M=200):

        #K: number of slices per iteration. The same K in max K-SWD.
        #M: number of spline knots of rational quadratic spline

        super().__init__()
        self.ndim = ndim
        if K is None:
            self.K = ndim
        else:
            self.K = K
        self.M = M

        ATi = torch.randn(self.ndim, self.K)
        Q, R = torch.linalg.qr(ATi)
        L = torch.sign(torch.diag(R))
        A = (Q * L)

        self.A = nn.Parameter(A)
        self.transform1D = RQspline(self.K, M)
//  [python]
def qr(x):  # pragma: no cover
    # inspired by
    # https://discuss.pytorch.org/t/multidimensional-svd/4366/2
    # prolonged here:
    if x.dim() == 2:
        result = torch.qr(x)
    else:
        batches = x.shape[:-2]
        other = x.shape[-2:]
        flat = x.view((-1,) + other)
        slices = flat.unbind(0)
        Q, R = [], []
        # I wish I had a parallel_for
        for i in range(flat.shape[0]):
            q, r = torch.qr(slices[i])
            Q += [q]
            R += [r]
        Q = torch.stack(Q).view(batches + Q[0].shape)
        R = torch.stack(R).view(batches + R[0].shape)
        result = Q, R
    return result
//  [python]
def qr(x):  # pragma: no cover
    # inspired by
    # https://discuss.pytorch.org/t/multidimensional-svd/4366/2
    # prolonged here:
    if x.dim() == 2:
        result = torch.linalg.qr(x)
    else:
        batches = x.shape[:-2]
        other = x.shape[-2:]
        flat = x.view((-1,) + other)
        slices = flat.unbind(0)
        Q, R = [], []
        # I wish I had a parallel_for
        for i in range(flat.shape[0]):
            q, r = torch.linalg.qr(slices[i])
            Q += [q]
            R += [r]
        Q = torch.stack(Q).view(batches + Q[0].shape)
        R = torch.stack(R).view(batches + R[0].shape)
        result = Q, R
    return result
//  [python]
def __init__(self, ndim, n_class, K=None, M=100):

        super().__init__()
        self.ndim = ndim
        self.n_class = n_class
        if K is None:
            self.K = ndim
        else:
            self.K = K
        self.M = M

        ATi = torch.randn(self.ndim, self.K)
        Q, R = torch.qr(ATi)
        L = torch.sign(torch.diag(R))
        A = (Q * L)

        self.A = nn.Parameter(A)
        self.transform1D = nn.ModuleList([RQspline(self.K, M) for i in range(self.n_class)])
//  [python]
def __init__(self, ndim, n_class, K=None, M=100):

        super().__init__()
        self.ndim = ndim
        self.n_class = n_class
        if K is None:
            self.K = ndim
        else:
            self.K = K
        self.M = M

        ATi = torch.randn(self.ndim, self.K)
        Q, R = torch.linalg.qr(ATi)
        L = torch.sign(torch.diag(R))
        A = (Q * L)

        self.A = nn.Parameter(A)
        self.transform1D = nn.ModuleList([RQspline(self.K, M) for i in range(self.n_class)])
//  [python]
def fit_A(self, data, sample='gaussian', MSWD_p=2, MSWD_max_iter=200, verbose=True):

        #fit the directions to apply 1D transform

        if verbose:
            tstart = start_timing(self.A.device)

        A, SWD = maxKSWDdirection(data, x2=sample, K=self.K, maxiter=MSWD_max_iter, p=MSWD_p)
        with torch.no_grad():
            SWD, indices = torch.sort(SWD, descending=True)
            A = A[:,indices]
            self.A[:] = torch.qr(A)[0] 

        if verbose:
            t = end_timing(tstart, self.A.device)
            print ('Fit A:', 'Time:', t, 'Wasserstein Distance:', SWD.tolist())
        return self
//  [python]
def fit_A(self, data, sample='gaussian', MSWD_p=2, MSWD_max_iter=200, verbose=True):

        #fit the directions to apply 1D transform

        if verbose:
            tstart = start_timing(self.A.device)

        A, SWD = maxKSWDdirection(data, x2=sample, K=self.K, maxiter=MSWD_max_iter, p=MSWD_p)
        with torch.no_grad():
            SWD, indices = torch.sort(SWD, descending=True)
            A = A[:,indices]
            self.A[:] = torch.linalg.qr(A)[0] 

        if verbose:
            t = end_timing(tstart, self.A.device)
            print ('Fit A:', 'Time:', t, 'Wasserstein Distance:', SWD.tolist())
        return self
//  [python]
def cpu_op_exec(self, input1, some):
        q, r = torch.qr(input1, some)
        return q.numpy(), r.numpy()
//  [python]
def cpu_op_exec(self, input1, some):
        q, r = torch.linalg.qr(input1, some)
        return q.numpy(), r.numpy()
//  [python]
def test_symsqrt():
    u.seed_random(1)
    torch.set_default_dtype(torch.float32)

    mat = torch.reshape(torch.arange(9) + 1, (3, 3)).float() + torch.eye(3) * 5
    mat = mat + mat.t()  # make symmetric
    smat = u.symsqrt(mat)
    u.check_close(mat, smat @ smat.t())
    u.check_close(mat, smat @ smat)

    def randomly_rotate(X):
        """Randomly rotate d,n data matrix X"""
        d, n = X.shape
        z = torch.randn((d, d), dtype=X.dtype)
        q, r = torch.qr(z)
        d = torch.diag(r)
        ph = d / abs(d)
        rot_mat = q * ph
        return rot_mat @ X

    n = 20
    d = 10
    X = torch.randn((d, n))

    # embed in a larger space
    X = torch.cat([X, torch.zeros_like(X)])
    X = randomly_rotate(X)
    cov = X @ X.t() / n
    sqrt, rank = u.symsqrt(cov, return_rank=True)
    assert rank == d
    assert torch.allclose(sqrt @ sqrt, cov, atol=1e-5)

    Y = torch.randn((d, n))
    Y = torch.cat([Y, torch.zeros_like(X)])
    Y = randomly_rotate(X)
    cov = u.Kron(X @ X.t(), Y @ Y.t())
    sqrt, rank = cov.symsqrt(return_rank=True)
    assert rank == d * d
    u.check_close(sqrt @ sqrt, cov, rtol=1e-4)

    X = torch.tensor([[7., 0, 0, 0, 0]]).t()
    X = randomly_rotate(X)
    cov = X @ X.t()
    u.check_close(u.sym_l2_norm(cov), 7 * 7)

    Y = torch.tensor([[8., 0, 0, 0, 0]]).t()
    Y = randomly_rotate(Y)
    cov = u.Kron(X @ X.t(), Y @ Y.t())
    u.check_close(cov.sym_l2_norm(), 7 * 7 * 8 * 8)
//  [python]
def test_symsqrt():
    u.seed_random(1)
    torch.set_default_dtype(torch.float32)

    mat = torch.reshape(torch.arange(9) + 1, (3, 3)).float() + torch.eye(3) * 5
    mat = mat + mat.t()  # make symmetric
    smat = u.symsqrt(mat)
    u.check_close(mat, smat @ smat.t())
    u.check_close(mat, smat @ smat)

    def randomly_rotate(X):
        """Randomly rotate d,n data matrix X"""
        d, n = X.shape
        z = torch.randn((d, d), dtype=X.dtype)
        q, r = torch.linalg.qr(z)
        d = torch.diag(r)
        ph = d / abs(d)
        rot_mat = q * ph
        return rot_mat @ X

    n = 20
    d = 10
    X = torch.randn((d, n))

    # embed in a larger space
    X = torch.cat([X, torch.zeros_like(X)])
    X = randomly_rotate(X)
    cov = X @ X.t() / n
    sqrt, rank = u.symsqrt(cov, return_rank=True)
    assert rank == d
    assert torch.allclose(sqrt @ sqrt, cov, atol=1e-5)

    Y = torch.randn((d, n))
    Y = torch.cat([Y, torch.zeros_like(X)])
    Y = randomly_rotate(X)
    cov = u.Kron(X @ X.t(), Y @ Y.t())
    sqrt, rank = cov.symsqrt(return_rank=True)
    assert rank == d * d
    u.check_close(sqrt @ sqrt, cov, rtol=1e-4)

    X = torch.tensor([[7., 0, 0, 0, 0]]).t()
    X = randomly_rotate(X)
    cov = X @ X.t()
    u.check_close(u.sym_l2_norm(cov), 7 * 7)

    Y = torch.tensor([[8., 0, 0, 0, 0]]).t()
    Y = randomly_rotate(Y)
    cov = u.Kron(X @ X.t(), Y @ Y.t())
    u.check_close(cov.sym_l2_norm(), 7 * 7 * 8 * 8)
//  [python]
def fit_A(self, data, sample='gaussian', weight=None, ndata_A=None, MSWD_p=2, MSWD_max_iter=200, pool=None, verbose=True):

        #fit the directions A to apply 1D transform

        if verbose:
            tstart = start_timing(self.A.device)

        if ndata_A is None or ndata_A > len(data):
            ndata_A = len(data)
        if sample != 'gaussian':
            if ndata_A > len(sample):
                ndata_A = len(sample)
            if ndata_A == len(sample):
                sample = sample.to(self.A.device)
            else:
                sample = sample[torch.randperm(len(sample), device=sample.device)[:ndata_A]].to(self.A.device)
        if ndata_A == len(data):
            data = data.to(self.A.device)
            if weight is not None:
                weight = weight.to(self.A.device)
        else:
            order = torch.randperm(len(data), device=data.device)[:ndata_A]
            data = data[order].to(self.A.device)
            if weight is not None:
                weight = weight[order].to(self.A.device)
        if weight is not None:
            weight = weight / torch.sum(weight)
            select = weight > 0
            data = data[select]
            weight = weight[select]

        A, SWD = maxKSWDdirection(data, x2=sample, weight=weight, K=self.K, maxiter=MSWD_max_iter, p=MSWD_p)
        with torch.no_grad():
            SWD, indices = torch.sort(SWD, descending=True)
            A = A[:,indices]
            self.A[:] = torch.qr(A)[0] 

        if verbose:
            t = end_timing(tstart, self.A.device)
            print ('Fit A:', 'Time:', t, 'Wasserstein Distance:', SWD.tolist())
        return self
//  [python]
def fit_A(self, data, sample='gaussian', weight=None, ndata_A=None, MSWD_p=2, MSWD_max_iter=200, pool=None, verbose=True):

        #fit the directions A to apply 1D transform

        if verbose:
            tstart = start_timing(self.A.device)

        if ndata_A is None or ndata_A > len(data):
            ndata_A = len(data)
        if sample != 'gaussian':
            if ndata_A > len(sample):
                ndata_A = len(sample)
            if ndata_A == len(sample):
                sample = sample.to(self.A.device)
            else:
                sample = sample[torch.randperm(len(sample), device=sample.device)[:ndata_A]].to(self.A.device)
        if ndata_A == len(data):
            data = data.to(self.A.device)
            if weight is not None:
                weight = weight.to(self.A.device)
        else:
            order = torch.randperm(len(data), device=data.device)[:ndata_A]
            data = data[order].to(self.A.device)
            if weight is not None:
                weight = weight[order].to(self.A.device)
        if weight is not None:
            weight = weight / torch.sum(weight)
            select = weight > 0
            data = data[select]
            weight = weight[select]

        A, SWD = maxKSWDdirection(data, x2=sample, weight=weight, K=self.K, maxiter=MSWD_max_iter, p=MSWD_p)
        with torch.no_grad():
            SWD, indices = torch.sort(SWD, descending=True)
            A = A[:,indices]
            self.A[:] = torch.linalg.qr(A)[0] 

        if verbose:
            t = end_timing(tstart, self.A.device)
            print ('Fit A:', 'Time:', t, 'Wasserstein Distance:', SWD.tolist())
        return self
//  [python]
def projected_process_auxiliary_matrices_I(K, M, Y, sigma):
    """
    If "x", "m" indicate the data and inducing points, and "k(.,.)" is
    the kernel function, then
    K = k(x, m)
    M = k(m, m)
    Y are the data values and sigma^2 is the ridge factor that will be
    added to the diagonal of the approximated covariance matrix.
    In projected process (PP) approximation the predictive distribution
    for a given test point "t" is:
    p(y|t,x,Y,m) = Normal(A @ mu, B - A @ nu @ A.t())
    where mu and nu which are independent from the test inputs are
    calculated here but A = k(t, m) and B = k(t, t) should be
    calculated outside this routine.
    """
    assert type(sigma) == float or sigma.numel() == 1
    # mu
    L, _ = jitcholesky(M)
    A = torch.cat([K, sigma * L.t()], dim=0)
    _Y = torch.cat([Y, torch.zeros(L.size(0))], dim=0)
    Q, R = torch.qr(A)
    mu = R.inverse() @ Q.t() @ _Y
    # nu
    i = L.inverse()
    B = K @ i.t()
    I = torch.eye(i.size(0))
    T = torch.linalg.cholesky(B.t() @ B / sigma**2 + I)
    nu = i.t() @ (I - torch.cholesky_inverse(T)) @ i
    return mu, nu
//  [python]
def projected_process_auxiliary_matrices_I(K, M, Y, sigma):
    """
    If "x", "m" indicate the data and inducing points, and "k(.,.)" is
    the kernel function, then
    K = k(x, m)
    M = k(m, m)
    Y are the data values and sigma^2 is the ridge factor that will be
    added to the diagonal of the approximated covariance matrix.
    In projected process (PP) approximation the predictive distribution
    for a given test point "t" is:
    p(y|t,x,Y,m) = Normal(A @ mu, B - A @ nu @ A.t())
    where mu and nu which are independent from the test inputs are
    calculated here but A = k(t, m) and B = k(t, t) should be
    calculated outside this routine.
    """
    assert type(sigma) == float or sigma.numel() == 1
    # mu
    L, _ = jitcholesky(M)
    A = torch.cat([K, sigma * L.t()], dim=0)
    _Y = torch.cat([Y, torch.zeros(L.size(0))], dim=0)
    Q, R = torch.linalg.qr(A)
    mu = R.inverse() @ Q.t() @ _Y
    # nu
    i = L.inverse()
    B = K @ i.t()
    I = torch.eye(i.size(0))
    T = torch.linalg.cholesky(B.t() @ B / sigma**2 + I)
    nu = i.t() @ (I - torch.cholesky_inverse(T)) @ i
    return mu, nu
//  [python]
def cpu_op_exec_tuple(self, input1, some):
        out = torch.qr(input1, some)
        output_q = out.Q
        output_r = out.R
        output_q = output_q.numpy()
        output_r = output_r.numpy()
        return output_q, output_r, out
//  [python]
def cpu_op_exec_tuple(self, input1, some):
        out = torch.linalg.qr(input1, some)
        output_q = out.Q
        output_r = out.R
        output_q = output_q.numpy()
        output_r = output_r.numpy()
        return output_q, output_r, out
//  [python]
def npu_op_exec(self, input1, some):
        q, r = torch.qr(input1, some)
        qout = q.to("cpu").numpy()
        rout = r.to("cpu").numpy()
        return qout, rout
//  [python]
def npu_op_exec(self, input1, some):
        q, r = torch.linalg.qr(input1, some)
        qout = q.to("cpu").numpy()
        rout = r.to("cpu").numpy()
        return qout, rout
//  [python]
def __init__(self, shape=[28,28,1], kernel=[4,4,1], shift=[0,0], K=None, M=200):

        assert shift[0] >= 0 and shift[0] < shape[0]
        assert shift[1] >= 0 and shift[1] < shape[1]
        assert len(shape) == 3 and len(kernel) == 3 and len(shift) == 2
        assert (kernel[0] <= shape[0]) and (kernel[1] <= shape[1])
        if shape[-1] == 1:
            assert kernel[-1] == 1
        else:
            assert (kernel[-1] == 1) or (kernel[-1] == shape[-1])

        super().__init__()
        self.register_buffer('shape', torch.tensor(shape))
        self.register_buffer('kernel', torch.tensor(kernel)) 
        self.register_buffer('shift', torch.tensor(shift))
        
        self.ndim_sub = (self.kernel[0]*self.kernel[1]*self.kernel[2]).item()

        if K is None:
            self.K = self.ndim_sub 
        else:
            self.K = K
            assert K <= self.ndim_sub
        self.M = M
        
        self.Nkernel_H = (self.shape[0] // self.kernel[0]).item()
        self.Nkernel_W = (self.shape[1] // self.kernel[1]).item()
        self.Nkernel_C = (self.shape[2] // self.kernel[2]).item()
        self.Nkernel = self.Nkernel_H * self.Nkernel_W * self.Nkernel_C 

        A = torch.zeros(self.Nkernel, self.ndim_sub, self.K)
        for i in range(self.Nkernel):
            ATi = torch.randn(self.ndim_sub, self.K)
            Q, R = torch.qr(ATi)
            L = torch.sign(torch.diag(R))
            A[i] = (Q * L)

        self.A = nn.Parameter(A)
        self.transform1D = RQspline(self.Nkernel*self.K, M)
//  [python]
def __init__(self, shape=[28,28,1], kernel=[4,4,1], shift=[0,0], K=None, M=200):

        assert shift[0] >= 0 and shift[0] < shape[0]
        assert shift[1] >= 0 and shift[1] < shape[1]
        assert len(shape) == 3 and len(kernel) == 3 and len(shift) == 2
        assert (kernel[0] <= shape[0]) and (kernel[1] <= shape[1])
        if shape[-1] == 1:
            assert kernel[-1] == 1
        else:
            assert (kernel[-1] == 1) or (kernel[-1] == shape[-1])

        super().__init__()
        self.register_buffer('shape', torch.tensor(shape))
        self.register_buffer('kernel', torch.tensor(kernel)) 
        self.register_buffer('shift', torch.tensor(shift))
        
        self.ndim_sub = (self.kernel[0]*self.kernel[1]*self.kernel[2]).item()

        if K is None:
            self.K = self.ndim_sub 
        else:
            self.K = K
            assert K <= self.ndim_sub
        self.M = M
        
        self.Nkernel_H = (self.shape[0] // self.kernel[0]).item()
        self.Nkernel_W = (self.shape[1] // self.kernel[1]).item()
        self.Nkernel_C = (self.shape[2] // self.kernel[2]).item()
        self.Nkernel = self.Nkernel_H * self.Nkernel_W * self.Nkernel_C 

        A = torch.zeros(self.Nkernel, self.ndim_sub, self.K)
        for i in range(self.Nkernel):
            ATi = torch.randn(self.ndim_sub, self.K)
            Q, R = torch.linalg.qr(ATi)
            L = torch.sign(torch.diag(R))
            A[i] = (Q * L)

        self.A = nn.Parameter(A)
        self.transform1D = RQspline(self.Nkernel*self.K, M)
//  [python]
def npu_op_exec_tuple(self, input1, some):
        out = torch.qr(input1.to("npu"), some)
        output_q = out.Q
        output_r = out.R
        output_q = output_q.to("cpu")
        output_r = output_r.to("cpu")
        output_q = output_q.numpy()
        output_r = output_r.numpy()
        return output_q, output_r, out
//  [python]
def npu_op_exec_tuple(self, input1, some):
        out = torch.linalg.qr(input1.to("npu"), some)
        output_q = out.Q
        output_r = out.R
        output_q = output_q.to("cpu")
        output_r = output_r.to("cpu")
        output_q = output_q.numpy()
        output_r = output_r.numpy()
        return output_q, output_r, out
//  [python]
def npu_op_exec_out(self, input1, some, input2, input3):
        torch.qr(input1, some, out=(input2, input3))
        qout = input2.to("cpu").numpy()
        rout = input3.to("cpu").numpy()
        return qout, rout
//  [python]
def npu_op_exec_out(self, input1, some, input2, input3):
        torch.linalg.qr(input1, mode='raw', out=(input2, input3))
        qout = input2.to("cpu").numpy()
        rout = input3.to("cpu").numpy()
        return qout, rout
//  [python]
def lstq(A, Y, lamb=0.0):
        """Differentiable least square."""
        q, r = torch.qr(A)
        return torch.inverse(r) @ q.permute(0, 2, 1) @ Y
//  [python]
def lstq(A, Y, lamb=0.0):
        """Differentiable least square."""
        q, r = torch.linalg.qr(A)
        return torch.inverse(r) @ q.permute(0, 2, 1) @ Y
//  [python]
def create_projection_matrix(m, d, seed=0):
    block_list = []
    current_seed = seed
    torch.manual_seed(current_seed)
    unstructured_block = torch.randn((d, d))
    q, _ = torch.qr(unstructured_block)
    q = torch.t(q)
    block_list.append(q[0:m])
    final_matrix = torch.vstack(block_list)

    current_seed += 1
    torch.manual_seed(current_seed)
    multiplier = torch.norm(torch.randn((m, d)), dim=1)

    return torch.matmul(torch.diag(multiplier), final_matrix)
//  [python]
def create_projection_matrix(m, d, seed=0):
    block_list = []
    current_seed = seed
    torch.manual_seed(current_seed)
    unstructured_block = torch.randn((d, d))
    q, _ = torch.linalg.qr(unstructured_block)
    q = torch.t(q)
    block_list.append(q[0:m])
    final_matrix = torch.vstack(block_list)

    current_seed += 1
    torch.manual_seed(current_seed)
    multiplier = torch.norm(torch.randn((m, d)), dim=1)

    return torch.matmul(torch.diag(multiplier), final_matrix)
//  [python]
def qr(
    x: torch.Tensor,
    /,
    *,
    mode: str = "reduced",
    out: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
) -> Tuple[torch.Tensor, torch.Tensor]:
    res = namedtuple("qr", ["Q", "R"])
    if mode == "reduced":
        q, r = torch.qr(x, some=True, out=out)
        ret = res(q, r)
    elif mode == "complete":
        q, r = torch.qr(x, some=False, out=out)
        ret = res(q, r)
    else:
        raise ivy.utils.exceptions.IvyException(
            "Only 'reduced' and 'complete' qr modes are allowed for the torch backend."
        )
    return ret
//  [python]
def qr(
    x: torch.Tensor,
    /,
    *,
    mode: str = "reduced",
    out: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
) -> Tuple[torch.Tensor, torch.Tensor]:
    res = namedtuple("qr", ["Q", "R"])
    if mode == "reduced":
        q, r = torch.linalg.qr(x, mode='reduced', out=out)
        ret = res(q, r)
    elif mode == "complete":
        q, r = torch.linalg.qr(x, mode='complete', out=out)
        ret = res(q, r)
    else:
        raise ivy.utils.exceptions.IvyException(
            "Only 'reduced' and 'complete' qr modes are allowed for the torch backend."
        )
    return ret
//  [python]
def make_linear_dataset(n, d, seed=0):
    """
    Create a dataset for training a deep linear network with n datapoints of dimension d.
    """
    torch.manual_seed(seed)
    X = (torch.qr(torch.randn(n, d))[0] * sqrt(n)).cuda()
    A = torch.randn(d, d).cuda()
    Y = X.mm(A.t())
    return TensorDataset(X, Y), TensorDataset(X, Y)
//  [python]
def make_linear_dataset(n, d, seed=0):
    """
    Create a dataset for training a deep linear network with n datapoints of dimension d.
    """
    torch.manual_seed(seed)
    X = (torch.linalg.qr(torch.randn(n, d))[0] * sqrt(n)).cuda()
    A = torch.randn(d, d).cuda()
    Y = X.mm(A.t())
    return TensorDataset(X, Y), TensorDataset(X, Y)
//  [python]
def orthogonal_matrix_chunk(cols, device=None):
    unstructured_block = torch.randn((cols, cols), device=device)
    q, r = torch.qr(unstructured_block.cpu(), some=True)
    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())
    q, r = map(lambda t: t.to(device), (q, r))
    return q.t()
//  [python]
def __init__(self, c):
        super(Invertible1x1ConvLUS, self).__init__()
        # Sample a random orthonormal matrix to initialize weights
        W, _ = torch.qr(torch.randn(c, c))
        # Ensure determinant is 1.0 not -1.0
        if torch.det(W) < 0:
            W[:, 0] = -1*W[:, 0]
        p, lower, upper = torch.lu_unpack(*torch.lu(W))

        self.register_buffer('p', p)
        # diagonals of lower will always be 1s anyway
        lower = torch.tril(lower, -1)
        lower_diag = torch.diag(torch.eye(c, c))
        self.register_buffer('lower_diag', lower_diag)
        self.lower = nn.Parameter(lower)
        self.upper_diag = nn.Parameter(torch.diag(upper))
        self.upper = nn.Parameter(torch.triu(upper, 1))
//  [python]
def __init__(self, c):
        super(Invertible1x1ConvLUS, self).__init__()
        # Sample a random orthonormal matrix to initialize weights
        W, _ = torch.linalg.qr(torch.randn(c, c))
        # Ensure determinant is 1.0 not -1.0
        if torch.det(W) < 0:
            W[:, 0] = -1*W[:, 0]
        p, lower, upper = torch.lu_unpack(*torch.lu(W))

        self.register_buffer('p', p)
        # diagonals of lower will always be 1s anyway
        lower = torch.tril(lower, -1)
        lower_diag = torch.diag(torch.eye(c, c))
        self.register_buffer('lower_diag', lower_diag)
        self.lower = nn.Parameter(lower)
        self.upper_diag = nn.Parameter(torch.diag(upper))
        self.upper = nn.Parameter(torch.triu(upper, 1))
//  [python]
def QRreduction(datas):
    
    ndatas = torch.qr(datas.permute(0,2,1)).R
    ndatas = ndatas.permute(0,2,1)
    return ndatas
//  [python]
def QRreduction(datas):
    
    ndatas = torch.linalg.qr(datas.permute(0,2,1)).R
    ndatas = ndatas.permute(0,2,1)
    return ndatas
//  [python]
def svd_randomized(M, k=10, device='cpu'):
    # http://gregorygundersen.com/blog/2019/01/17/randomized-svd/
    # http://algorithm-interest-group.me/assets/slides/randomized_SVD.pdf
    n_samples, n_dims = M.shape[0], M.shape[1]
    P = torch.randn([n_dims, k]).float().to(device)
    M_P = torch.mm(M, P)
    Q, _ = torch.qr(M_P)
    B = torch.mm(Q.T, M)
    U, S, V = torch.svd(B)
    U = torch.mm(Q, U)

    return U, S, V
//  [python]
def svd_randomized(M, k=10, device='cpu'):
    n_samples, n_dims = M.shape[0], M.shape[1]
    P = torch.randn([n_dims, k]).float().to(device)
    M_P = torch.mm(M, P)
    Q, _ = torch.linalg.qr(M_P)
    B = torch.mm(Q.T, M)
    U, S, V = torch.svd(B)
    U = torch.mm(Q, U)

    return U, S, V
//  [python]
def projected(self):
        a = self.active
        L = self.K[a][:, a].cholesky()
        sigma = self.sigma
        A = torch.cat([self.K[:, a] / sigma.view(-1, 1), L.t()])
        O = torch.zeros(L.size(0)).type(L.type())
        Y = torch.cat([self.y / sigma, O])
        Q, R = torch.qr(A)
        mu = R.inverse() @ Q.t() @ Y
        delta = ((self.K[:, a] @ mu).view(-1) - self.y).abs()
        return a, mu, delta.mean(), delta.max()
//  [python]
def projected(self):
        a = self.active
        L = self.K[a][:, a].cholesky()
        sigma = self.sigma
        A = torch.cat([self.K[:, a] / sigma.view(-1, 1), L.t()])
        O = torch.zeros(L.size(0)).type(L.type())
        Y = torch.cat([self.y / sigma, O])
        Q, R = torch.linalg.qr(A)
        mu = R.inverse() @ Q.t() @ Y
        delta = ((self.K[:, a] @ mu).view(-1) - self.y).abs()
        return a, mu, delta.mean(), delta.max()
//  [python]
def qr(
    torch: Any,
    tensor: Tensor,
    pivot_axis: int,
    non_negative_diagonal: bool = False
) -> Tuple[Tensor, Tensor]:
  """Computes the QR decomposition of a tensor.

  The QR decomposition is performed by treating the tensor as a matrix,
  with an effective left (row) index resulting from combining the axes
  `tensor.shape[:pivot_axis]` and an effective right (column) index
  resulting from combining the axes `tensor.shape[pivot_axis:]`.

  For example, if `tensor` had a shape (2, 3, 4, 5) and `pivot_axis` was 2,
  then `q` would have shape (2, 3, 6), and `r` would have shape (6, 4, 5).

  The output consists of two tensors `Q, R` such that:
  
    Q[i1,...,iN, j] * R[j, k1,...,kM] == tensor[i1,...,iN, k1,...,kM]
  
  `R` is an upper triangular matrix, `Q` is an orthonormal matrix
  Note that the output ordering matches numpy.linalg.svd rather than tf.svd.

  Args:
    tf: The tensorflow module.
    tensor: A tensor to be decomposed.
    pivot_axis: Where to split the tensor's axes before flattening into a
      matrix.

  Returns:
    Q: Left tensor factor.
    R: Right tensor factor.
  """

  left_dims = list(tensor.shape)[:pivot_axis]
  right_dims = list(tensor.shape)[pivot_axis:]

  tensor = torch.reshape(tensor, (np.prod(left_dims), np.prod(right_dims)))
  q, r = torch.linalg.qr(tensor)
  if non_negative_diagonal:
    phases = torch.sign(torch.diagonal(r))
    q = q * phases
    r = phases[:, None] * r
  center_dim = q.shape[1]
  q = torch.reshape(q, list(left_dims) + [center_dim])
  r = torch.reshape(r, [center_dim] + list(right_dims))
  return q, r
//  [python]
def __init__(self, in_channel):
        super().__init__()
        w_shape = [in_channel, in_channel]
        # Process of sampling random rotation matrix
        w_init, _ = torch.qr(torch.randn(*w_shape))
        P, L, U = torch.lu_unpack(*torch.lu(w_init))
        s = torch.diag(U)

        self.register_buffer("L_mask", torch.tril(torch.ones(w_shape), -1))
        self.register_buffer("eye", torch.eye(*w_shape))
        self.register_buffer("P", P)
        self.register_buffer("s_sign", torch.sign(s))

        self.log_s = nn.Parameter(torch.log(torch.abs(s)))
        self.L = nn.Parameter(L)
        self.U = nn.Parameter(U)
//  [python]
def __init__(self, in_channel):
        super().__init__()
        w_shape = [in_channel, in_channel]
        # Process of sampling random rotation matrix
        w_init, _ = torch.linalg.qr(torch.randn(*w_shape))
        P, L, U = torch.lu_unpack(*torch.lu(w_init))
        s = torch.diag(U)

        self.register_buffer("L_mask", torch.tril(torch.ones(w_shape), -1))
        self.register_buffer("eye", torch.eye(*w_shape))
        self.register_buffer("P", P)
        self.register_buffer("s_sign", torch.sign(s))

        self.log_s = nn.Parameter(torch.log(torch.abs(s)))
        self.L = nn.Parameter(L)
        self.U = nn.Parameter(U)
//  [python]
def random_rotation():
    q, r = torch.qr(torch.randn(3, 3))
    return q
//  [python]
def random_rotation():
    q, r = torch.linalg.qr(torch.randn(3, 3))
    return q