jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def lstq(A, F_u, F_v, lamda=0.01):
    # cols = A.shape[2]
    # assert all(cols == torch.linalg.matrix_rank(A))  # something better?
    try:
        Q, R = torch.linalg.qr(A)
        theta_x = torch.bmm(torch.bmm(torch.linalg.inv(R), Q.transpose(1, 2)), F_u)
        theta_y = torch.bmm(torch.bmm(torch.linalg.inv(R), Q.transpose(1, 2)), F_v)
    except:
        LOGGER.exception("Least Squares failed")
        sys.exit(-1)
    return theta_x, theta_y
//  [python]
def parallel_orthogonal_(tensor, gain=1):
    if tensor.ndimension() < 3:
        raise ValueError("Only tensors with 3 or more dimensions are supported")

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

    qs = []
    for flat_tensor in torch.unbind(flattened, dim=0):
        if rows < cols:
            flat_tensor.t_()

        # Compute the qr factorization
        q, r = torch.qr(flat_tensor)
        # 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_()
        qs.append(q)

    qs = torch.stack(qs, dim=0)

    with torch.no_grad():
        tensor.view_as(qs).copy_(qs)
        tensor.mul_(gain)
    return tensor
//  [python]
def parallel_orthogonal_(tensor, gain=1):
    if tensor.ndimension() < 3:
        raise ValueError("Only tensors with 3 or more dimensions are supported")

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

    qs = []
    for flat_tensor in torch.unbind(flattened, dim=0):
        if rows < cols:
            flat_tensor.t_()

        # Compute the qr factorization
        q, r = torch.linalg.qr(flat_tensor)
        # 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_()
        qs.append(q)

    qs = torch.stack(qs, dim=0)

    with torch.no_grad():
        tensor.view_as(qs).copy_(qs)
        tensor.mul_(gain)
    return tensor
//  [python]
def sliced_wasserstein_sphere(Xs, Xt, num_projections, device, u_weights=None, v_weights=None, p=2):
    """
        Compute the sliced-Wasserstein distance on the sphere.

        Parameters:
        Xs: ndarray, shape (n_samples_u, dim)
            Samples in the source domain
        Xt: ndarray, shape (n_samples_v, dim)
            Samples in the target domain
        num_projections: int
            Number of projections
        device: str
        p: float
            Power of SW. Need to be >= 1.
    """
    d = Xs.shape[1]
    
    ## Uniforms and independent samples on the Stiefel manifold V_{d,2}
    Z = torch.randn((num_projections,d,2), device=device)
    Q, R = torch.qr(Z)
        
    return sliced_cost(Xs, Xt, Q, p=p, u_weights=u_weights, v_weights=v_weights)
//  [python]
def sliced_wasserstein_sphere(Xs, Xt, num_projections, device, u_weights=None, v_weights=None, p=2):
    """
        Compute the sliced-Wasserstein distance on the sphere.

        Parameters:
        Xs: ndarray, shape (n_samples_u, dim)
            Samples in the source domain
        Xt: ndarray, shape (n_samples_v, dim)
            Samples in the target domain
        num_projections: int
            Number of projections
        device: str
        p: float
            Power of SW. Need to be >= 1.
    """
    d = Xs.shape[1]
    
    ## Uniforms and independent samples on the Stiefel manifold V_{d,2}
    Z = torch.randn((num_projections,d,2), device=device)
    U, _ = torch.linalg.qr(Z)
        
    return sliced_cost(Xs, Xt, U, p=p, u_weights=u_weights, v_weights=v_weights)
//  [python]
def _orthogonal_matrix(dim: int) -> Tensor:
    r"""Get an orthogonal matrix by applying QR decomposition."""
    # Random matrix from normal distribution
    mat = torch.randn((dim, dim))
    # QR decomposition to two orthogonal matrices
    q, _ = torch.qr(mat.cpu())
    return q.t()
//  [python]
def _orthogonal_matrix(dim: int) -> Tensor:
    r"""Get an orthogonal matrix by applying QR decomposition."""
    # Random matrix from normal distribution
    mat = torch.randn((dim, dim))
    # QR decomposition to two orthogonal matrices
    q, _ = torch.linalg.qr(mat.cpu(), mode='reduced')
    return q.t()
//  [python]
def sample_base(cls, batch_size, dim=None):
        z = torch.randn(batch_size, 3, 3, dtype=torch.cfloat)
        q, r = torch.qr(z)
        d = torch.diagonal(r, dim1=-2, dim2=-1)
        ph = d / d.abs()
        q = q @ torch.diag_embed(ph) @ q
        q = q / torch.det(q).pow(1/3)[..., None, None]
        return q
//  [python]
def sample_base(cls, batch_size, dim=None):
        z = torch.randn(batch_size, 3, 3, dtype=torch.cfloat)
        q, r = torch.linalg.qr(z)
        d = torch.diagonal(r, dim1=-2, dim2=-1)
        ph = d / d.abs()
        q = q @ torch.diag_embed(ph) @ q
        q = q / torch.det(q).pow(1/3)[..., None, None]
        return q
//  [python]
def sliced_wasserstein_sphere_unif(Xs, num_projections, device):
    """
        Compute the SSW2 on the sphere w.r.t. a uniform distribution.

        Parameters:
        Xs: ndarray, shape (n_samples_u, dim)
            Samples in the source domain
        num_projections: int
            Number of projections
        device: str.
    """
    n, d = Xs.shape
    
    ## Uniforms and independent samples on the Stiefel manifold V_{d,2}
    Z = torch.randn((num_projections,d,2), device=device)
    U, _ = torch.qr(Z)
    
    ## Projection on S^1
    ## Projection on plane
    Xps = torch.matmul(torch.transpose(U,1,2)[:,None], Xs[:,:,None]).reshape(num_projections, n, 2)
    ## Projection on sphere
    Xps = F.normalize(Xps, p=2, dim=-1)
    ## Get coords
    Xps = (torch.atan2(-Xps[:,:,1], -Xps[:,:,0])+np.pi)/(2*np.pi)
        
    return torch.mean(w2_unif_circle(Xps))
//  [python]
def sliced_wasserstein_sphere_unif(Xs, num_projections, device):
    """
        Compute the SSW2 on the sphere w.r.t. a uniform distribution.

        Parameters:
        Xs: ndarray, shape (n_samples_u, dim)
            Samples in the source domain
        num_projections: int
            Number of projections
        device: str.
    """
    n, d = Xs.shape
    
    ## Uniforms and independent samples on the Stiefel manifold V_{d,2}
    Z = torch.randn((num_projections,d,2), device=device)
    U, _ = torch.linalg.qr(Z)
    
    ## Projection on S^1
    ## Projection on plane
    Xps = torch.matmul(torch.transpose(U,1,2)[:,None], Xs[:,:,None]).reshape(num_projections, n, 2)
    ## Projection on sphere
    Xps = F.normalize(Xps, p=2, dim=-1)
    ## Get coords
    Xps = (torch.atan2(-Xps[:,:,1], -Xps[:,:,0])+np.pi)/(2*np.pi)
        
    return torch.mean(w2_unif_circle(Xps))
//  [python]
def test_orgqr(self, device, dtype):
        def generate_reflectors_and_tau(A):
            """
            This function uses numpy.linalg.qr with mode "raw" to extract output of LAPACK's geqrf.
            There is torch.geqrf function but it doesn't work with complex-valued input.
            """
            if A.numel() > 0:
                A_cpu = A.cpu()
                flattened_batch_shape = [-1, *A_cpu.shape[-2:]]
                reflectors = torch.empty_like(A_cpu).view(*flattened_batch_shape)
                tau_shape = [*A_cpu.shape[:-2], A_cpu.shape[-1]]
                tau = torch.empty(tau_shape, dtype=dtype).view(-1, A_cpu.shape[-1])
                for A_i, reflectors_i, tau_i in zip(A_cpu.contiguous().view(*flattened_batch_shape), reflectors, tau):
                    reflectors_tmp, tau_i[:] = map(torch.from_numpy, np.linalg.qr(A_i, mode='raw'))
                    reflectors_i[:] = reflectors_tmp.T
                reflectors = reflectors.view(*A_cpu.shape)
                tau = tau.view(tau_shape)
                return reflectors.to(A.device), tau.to(A.device)

            reflectors = torch.empty_like(A)
            tau = torch.empty(*A.shape[:-2], A.shape[-1], dtype=dtype, device=device)
            return reflectors, tau

        def run_test(shape):
            A = torch.randn(*shape, dtype=dtype, device=device)
            reflectors, tau = generate_reflectors_and_tau(A)
            Q, _ = torch.qr(A)
            actual = torch.orgqr(reflectors, tau)
            # torch.linalg.qr does not work correctly for zero batch dimension tensors
            # see https://github.com/pytorch/pytorch/issues/50576
            if (A.numel() > 0):
                self.assertEqual(Q, actual)
            else:
                self.assertTrue(actual.shape == shape)

            out = torch.empty_like(A)
            ans = torch.orgqr(reflectors, tau, out=out)
            self.assertEqual(ans, out)
            if (A.numel() > 0):
                self.assertEqual(Q, out)

        shapes = [(0, 0), (5, 0),  # Empty matrix
                  (5, 5), (5, 3),  # Single matrix
                  (0, 0, 0), (0, 5, 5), (0, 5, 3),  # Zero batch dimension tensors
                  (2, 5, 5), (2, 5, 3),  # 3-dim tensors
                  (2, 1, 5, 5), (2, 1, 5, 3)]  # 4-dim tensors
        for shape in shapes:
            run_test(shape)
//  [python]
def test_orgqr(self, device, dtype):
        def generate_reflectors_and_tau(A):
            """
            This function uses numpy.linalg.qr with mode "raw" to extract output of LAPACK's geqrf.
            There is torch.geqrf function but it doesn't work with complex-valued input.
            """
            if A.numel() > 0:
                A_cpu = A.cpu()
                flattened_batch_shape = [-1, *A_cpu.shape[-2:]]
                reflectors = torch.empty_like(A_cpu).view(*flattened_batch_shape)
                tau_shape = [*A_cpu.shape[:-2], A_cpu.shape[-1]]
                tau = torch.empty(tau_shape, dtype=dtype).view(-1, A_cpu.shape[-1])
                for A_i, reflectors_i, tau_i in zip(A_cpu.contiguous().view(*flattened_batch_shape), reflectors, tau):
                    reflectors_tmp, tau_i[:] = map(torch.from_numpy, np.linalg.qr(A_i, mode='raw'))
                    reflectors_i[:] = reflectors_tmp.T
                reflectors = reflectors.view(*A_cpu.shape)
                tau = tau.view(tau_shape)
                return reflectors.to(A.device), tau.to(A.device)

            reflectors = torch.empty_like(A)
            tau = torch.empty(*A.shape[:-2], A.shape[-1], dtype=dtype, device=device)
            return reflectors, tau

        def run_test(shape):
            A = torch.randn(*shape, dtype=dtype, device=device)
            reflectors, tau = generate_reflectors_and_tau(A)
            expected, _ = torch.linalg.qr(A)
            actual = torch.orgqr(reflectors, tau)
            # torch.linalg.qr does not work correctly for zero batch dimension tensors
            # see https://github.com/pytorch/pytorch/issues/50576
            if (A.numel() > 0):
                self.assertEqual(expected, actual)
            else:
                self.assertTrue(actual.shape == shape)

            out = torch.empty_like(A)
            ans = torch.orgqr(reflectors, tau, out=out)
            self.assertEqual(ans, out)
            if (A.numel() > 0):
                self.assertEqual(expected, out)

        shapes = [(0, 0), (5, 0),  # Empty matrix
                  (5, 5), (5, 3),  # Single matrix
                  (0, 0, 0), (0, 5, 5), (0, 5, 3),  # Zero batch dimension tensors
                  (2, 5, 5), (2, 5, 3),  # 3-dim tensors
                  (2, 1, 5, 5), (2, 1, 5, 3)]  # 4-dim tensors
        for shape in shapes:
            run_test(shape)
//  [python]
def test_linalg_qr() -> None:
    x: torch.Tensor[torch.float32, L[2], L[3], L[4]]
    y: Tuple[
        torch.Tensor[torch.float32, L[2], L[3], L[3]],
        torch.Tensor[torch.float32, L[2], L[3], L[4]],
    ] = torch.qr(x)
    # pyre-fixme[9]: y_error has type `Tuple[Tensor[torch.float32, typing_extensions....
    y_error: Tuple[
        torch.Tensor[torch.float32, L[2], L[3], L[99]],
        torch.Tensor[torch.float32, L[2], L[3], L[4]],
    ] = torch.qr(x)
    y2: Tuple[
        torch.Tensor[torch.float32, L[2], L[3], L[3]],
        torch.Tensor[torch.float32, L[2], L[3], L[4]],
    ] = torch.qr(x, mode="complete")
//  [python]
def test_linalg_qr() -> None:
    x: torch.Tensor[torch.float32, L[2], L[3], L[4]]
    y: Tuple[
        torch.Tensor[torch.float32, L[2], L[3], L[3]],
        torch.Tensor[torch.float32, L[2], L[3], L[4]],
    ] = torch.linalg.qr(x)
    # pyre-fixme[9]: y_error has type `Tuple[Tensor[torch.float32, typing_extensions....
    y_error: Tuple[
        torch.Tensor[torch.float32, L[2], L[3], L[99]],
        torch.Tensor[torch.float32, L[2], L[3], L[4]],
    ] = torch.linalg.qr(x)
    y2: Tuple[
        torch.Tensor[torch.float32, L[2], L[3], L[3]],
        torch.Tensor[torch.float32, L[2], L[3], L[4]],
    ] = torch.linalg.qr(x, mode="complete")
//  [python]
def random_so3(num):
    return torch.qr(torch.randn(num, 3, 3))[0]
//  [python]
def random_so3(num):
    return torch.linalg.qr(torch.randn(num, 3, 3)).Q
//  [python]
def true_poweriteration(input: torch.Tensor, loop, rank, p_base=None, q_base=None):
    # input size [batch,num_head,seq_len,model_dim/num_head]
    # -> [batch,seq_len,model_dim] -> [batch * seq_len,model_dim]
    # p_base = torch.rand(input.shape[3] * input.shape[1], rank).to(device)
    # q_base = torch.rand(input.shape[0] * input.shape[2], rank).to(device)
    batch, num_head, seq_len, sep_dim = input.shape
    input = (
        input.permute(0, 2, 1, 3).contiguous().view(batch, seq_len, sep_dim * num_head)
    )  # convert to 32bits for qr decomposition
    input = input.view(batch, seq_len, sep_dim * num_head)
    input = input.float()
    if q_base is not None and p_base is not None:
        p_base[0] = p_base[0].float()
        q_base[0] = q_base[0].float()
    else:
        p_base = [torch.rand(batch, sep_dim * num_head, rank).to(input.device).float()]
        q_base = [torch.rand(batch, seq_len, rank).to(input.device).float()]
    # 3 calculation = loop * (matmul) + 2 * qrO(n^2)
    for i in range(loop):
        if i == loop - 1:
            p_base[0] = torch.qr(p_base[0].float()).Q
        q_base[0] = input @ p_base[0]
        if i == loop - 1:
            q_base[0] = torch.qr(q_base[0].float()).Q
        p_base[0] = torch.transpose(input, 1, 2) @ q_base[0]
    # input = q_base[0] @ torch.transpose(p_base[0], 0, 1)
    # input = input.view(batch, seq_len, num_head, sep_dim)
    # input = input.permute(0, 2, 1, 3)
    # input = input.type(torch.bfloat16)
    p_base[0] = p_base[0].half()
    q_base[0] = q_base[0].half()
    return p_base, q_base
//  [python]
def true_poweriteration(input: torch.Tensor, loop, rank, p_base=None, q_base=None):
    # input size [batch,num_head,seq_len,model_dim/num_head]
    # -> [batch,seq_len,model_dim] -> [batch * seq_len,model_dim]
    # p_base = torch.rand(input.shape[3] * input.shape[1], rank).to(device)
    # q_base = torch.rand(input.shape[0] * input.shape[2], rank).to(device)
    batch, num_head, seq_len, sep_dim = input.shape
    input = (
        input.permute(0, 2, 1, 3).contiguous().view(batch, seq_len, sep_dim * num_head)
    )  # convert to 32bits for qr decomposition
    input = input.view(batch, seq_len, sep_dim * num_head)
    input = input.float()
    if q_base is not None and p_base is not None:
        p_base[0] = p_base[0].float()
        q_base[0] = q_base[0].float()
    else:
        p_base = [torch.rand(batch, sep_dim * num_head, rank).to(input.device).float()]
        q_base = [torch.rand(batch, seq_len, rank).to(input.device).float()]
    # 3 calculation = loop * (matmul) + 2 * qrO(n^2)
    for i in range(loop):
        if i == loop - 1:
            p_base[0] = torch.linalg.qr(p_base[0].float()).Q
        q_base[0] = input @ p_base[0]
        if i == loop - 1:
            q_base[0] = torch.linalg.qr(q_base[0].float()).Q
        p_base[0] = torch.transpose(input, 1, 2) @ q_base[0]
    # input = q_base[0] @ torch.transpose(p_base[0], 0, 1)
    # input = input.view(batch, seq_len, num_head, sep_dim)
    # input = input.permute(0, 2, 1, 3)
    # input = input.type(torch.bfloat16)
    p_base[0] = p_base[0].half()
    q_base[0] = q_base[0].half()
    return p_base, q_base
//  [python]
def test_equivariance1(self):
        torch.manual_seed(0)

        in_features = 3
        out_features = 4
        edge_features = 5
        hidden_features = 6
        num_layers = 7
        attention = False
        use_tanh = True
        egnn = EGNN(
            in_features,
            out_features,
            edge_features,
            hidden_features,
            num_layers,
            attention,
            use_tanh,
        )
        x = torch.randn(10, in_features)
        pos = torch.randn(10, 3)
        edge_index = torch.tensor(
            [[0, 1, 2, 3, 4, 5, 6, 7, 8, 9], [1, 2, 3, 4, 5, 6, 7, 8, 9, 0]]
        )
        mask = torch.tensor([True] * 10)
        edge_attr = torch.randn(10, edge_features)

        R, _ = torch.qr(torch.randn(3, 3))

        x1, pos1 = egnn(x, pos, edge_index, mask, edge_attr)
        x2, pos2 = egnn(x, pos @ R, edge_index, mask, edge_attr)
        self.assertTrue(torch.allclose(x1, x2))
        self.assertTrue(torch.allclose(pos1 @ R, pos2))
//  [python]
def test_equivariance1(self):
        torch.manual_seed(0)

        in_features = 3
        out_features = 4
        edge_features = 5
        hidden_features = 6
        num_layers = 7
        attention = False
        use_tanh = True
        egnn = EGNN(
            in_features,
            out_features,
            edge_features,
            hidden_features,
            num_layers,
            attention,
            use_tanh,
        )
        x = torch.randn(10, in_features)
        pos = torch.randn(10, 3)
        edge_index = torch.tensor(
            [[0, 1, 2, 3, 4, 5, 6, 7, 8, 9], [1, 2, 3, 4, 5, 6, 7, 8, 9, 0]]
        )
        mask = torch.tensor([True] * 10)
        edge_attr = torch.randn(10, edge_features)

        R, _ = torch.linalg.qr(torch.randn(3, 3))

        x1, pos1 = egnn(x, pos, edge_index, mask, edge_attr)
        x2, pos2 = egnn(x, pos @ R, edge_index, mask, edge_attr)
        self.assertTrue(torch.allclose(x1, x2))
        self.assertTrue(torch.allclose(pos1 @ R, pos2))
//  [python]
def __init__(self, D, d, params, names, seed=_DEFAULT_SEED):
        super().__init__(None, (D, d))
        self.info = (D, d, seed)
        self.P = torch.randn(D, d)
        Q, R = torch.qr(self.P)
        self.P = Q
//  [python]
def __init__(self, D, d, params, names, seed=_DEFAULT_SEED):
        super().__init__(None, (D, d))
        self.info = (D, d, seed)
        self.P = torch.randn(D, d)
        self.P, _ = torch.linalg.qr(self.P, mode="reduced")
//  [python]
def test_equivariance2(self):
        torch.manual_seed(1)

        in_features = 7
        out_features = 6
        edge_features = 5
        hidden_features = 4
        num_layers = 3
        attention = True
        use_tanh = False
        egnn = EGNN(
            in_features,
            out_features,
            edge_features,
            hidden_features,
            num_layers,
            attention,
            use_tanh,
        )
        x = torch.randn(10, in_features)
        pos = torch.randn(10, 3)
        edge_index = torch.tensor(
            [[0, 1, 2, 3, 4, 5, 6, 7, 8, 9], [1, 2, 3, 4, 5, 6, 7, 8, 9, 0]]
        )
        mask = torch.tensor([True] * 5 + [False] * 5)
        edge_attr = torch.randn(10, edge_features)

        R, _ = torch.qr(torch.randn(3, 3))

        x1, pos1 = egnn(x, pos, edge_index, mask, edge_attr)
        x2, pos2 = egnn(x, pos @ R, edge_index, mask, edge_attr)
        self.assertTrue(torch.allclose(x1, x2))
        self.assertTrue(torch.allclose(pos1 @ R, pos2))
//  [python]
def test_equivariance2(self):
        torch.manual_seed(1)

        in_features = 7
        out_features = 6
        edge_features = 5
        hidden_features = 4
        num_layers = 3
        attention = True
        use_tanh = False
        egnn = EGNN(
            in_features,
            out_features,
            edge_features,
            hidden_features,
            num_layers,
            attention,
            use_tanh,
        )
        x = torch.randn(10, in_features)
        pos = torch.randn(10, 3)
        edge_index = torch.tensor(
            [[0, 1, 2, 3, 4, 5, 6, 7, 8, 9], [1, 2, 3, 4, 5, 6, 7, 8, 9, 0]]
        )
        mask = torch.tensor([True] * 5 + [False] * 5)
        edge_attr = torch.randn(10, edge_features)

        R, _ = torch.linalg.qr(torch.randn(3, 3))

        x1, pos1 = egnn(x, pos, edge_index, mask, edge_attr)
        x2, pos2 = egnn(x, pos @ R, edge_index, mask, edge_attr)
        self.assertTrue(torch.allclose(x1, x2))
        self.assertTrue(torch.allclose(pos1 @ R, pos2))
//  [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")

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

    if rows < cols:
        flattened.t_()

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

    if rows < cols:
        q.t_()

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

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

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

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

    if rows < cols:
        flattened.t_()

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

    if rows < cols:
        q.t_()

    with torch.no_grad():
        tensor.view_as(q).copy_(q)
        tensor.mul_(gain)
    return tensor
//  [python]
def true_poweriteration_quantized(
    input: torch.Tensor, loop, rank, p_base=None, q_base=None
):
    # input size [batch,num_head,seq_len,model_dim/num_head]
    # -> [batch,seq_len,model_dim] -> [batch * seq_len,model_dim]
    # p_base = torch.rand(input.shape[3] * input.shape[1], rank).to(device)
    # q_base = torch.rand(input.shape[0] * input.shape[2], rank).to(device)
    batch, num_head, seq_len, sep_dim = input.shape
    input = (
        input.permute(0, 2, 1, 3).contiguous().view(batch, seq_len, sep_dim * num_head)
    )  # convert to 32bits for qr decomposition
    input = input.view(batch, seq_len, sep_dim * num_head)
    input = input.float()
    if q_base is not None and p_base is not None:
        p_base[0] = p_base[0].float()
        q_base[0] = q_base[0].float()
    else:
        p_base = [torch.rand(batch, sep_dim * num_head, rank).to(input.device).float()]
        q_base = [torch.rand(batch, seq_len, rank).to(input.device).float()]
    # 3 calculation = loop * (matmul) + 2 * qrO(n^2)
    for i in range(loop):
        if i == loop - 1:
            q, r = torch.qr(p_base[0].float())
            p_base[0] = q
        q_base[0] = input @ p_base[0]
        if i == loop - 1:
            q, r = torch.qr(q_base[0].float())
            q_base[0] = q
        p_base[0] = torch.transpose(input, 1, 2) @ q_base[0]
    # input = q_base[0] @ torch.transpose(p_base[0], 0, 1)
    # input = input.view(batch, seq_len, num_head, sep_dim)
    # input = input.permute(0, 2, 1, 3)
    # input = input.type(torch.bfloat16)

    # p_base[0] = p_base[0].half()
    # q_base[0] = q_base[0].half()
    #### compress p and q base to 8bits
    p_base[0], shape_p, min_p, scale_p = true_uniform_quantization_compress(
        p_base[0], 8
    )
    q_base[0], shape_q, min_q, scale_q = true_uniform_quantization_compress(
        q_base[0], 8
    )
    return p_base, q_base, shape_p, shape_q, min_p, min_q, scale_p, scale_q
//  [python]
def true_poweriteration_quantized(
    input: torch.Tensor, loop, rank, p_base=None, q_base=None
):
    # input size [batch,num_head,seq_len,model_dim/num_head]
    # -> [batch,seq_len,model_dim] -> [batch * seq_len,model_dim]
    # p_base = torch.rand(input.shape[3] * input.shape[1], rank).to(device)
    # q_base = torch.rand(input.shape[0] * input.shape[2], rank).to(device)
    batch, num_head, seq_len, sep_dim = input.shape
    input = (
        input.permute(0, 2, 1, 3).contiguous().view(batch, seq_len, sep_dim * num_head)
    )  # convert to 32bits for qr decomposition
    input = input.view(batch, seq_len, sep_dim * num_head)
    input = input.float()
    if q_base is not None and p_base is not None:
        p_base[0] = p_base[0].float()
        q_base[0] = q_base[0].float()
    else:
        p_base = [torch.rand(batch, sep_dim * num_head, rank).to(input.device).float()]
        q_base = [torch.rand(batch, seq_len, rank).to(input.device).float()]
    # 3 calculation = loop * (matmul) + 2 * qrO(n^2)
    for i in range(loop):
        if i == loop - 1:
            p_base[0] = torch.linalg.qr(p_base[0].float()).Q
        q_base[0] = input @ p_base[0]
        if i == loop - 1:
            q_base[0] = torch.linalg.qr(q_base[0].float()).Q
        p_base[0] = torch.transpose(input, 1, 2) @ q_base[0]
    # input = q_base[0] @ torch.transpose(p_base[0], 0, 1)
    # input = input.view(batch, seq_len, num_head, sep_dim)
    # input = input.permute(0, 2, 1, 3)
    # input = input.type(torch.bfloat16)

    # p_base[0] = p_base[0].half()
    # q_base[0] = q_base[0].half()
    #### compress p and q base to 8bits
    p_base[0], shape_p, min_p, scale_p = true_uniform_quantization_compress(
        p_base[0], 8
    )
    q_base[0], shape_q, min_q, scale_q = true_uniform_quantization_compress(
        q_base[0], 8
    )
    return p_base, q_base, shape_p, shape_q, min_p, min_q, scale_p, scale_q
//  [python]
def _get_random_ortho_matrix(
        blocks: int,
        dim: int,
        device: torch.device,
        norm_distribution: NormDistribution = NormDistribution.Uniform,
    ) -> torch.Tensor:
        r"""
        Generate a random matrix whose rows are exactly orthonormal

        "How to generate random matrices from the classical compact groups", Mezzadri, 2007
        https://arxiv.org/pdf/math-ph/0609050v2.pdf

        .. note: the typical qr decomposition does not give uniform results, qr decomposition is not
        unique and the qr decomposition routines are biased towards numerical stability. See the above
        paper for more information.

        .. note: this does not follow the original implementation from the Performers authors.
        see docs/assets/kde plots to visualize the impact of using the R signs to correct Q
        """

        H = torch.randn((blocks, dim, dim), device=device, requires_grad=False)

        # Randomly scale the norms of the features, Xi distributed
        if norm_distribution == NormDistribution.Xi:
            # NOTE: This averages to sqrt(d)
            norms = torch.sqrt(torch.einsum("...d,...d->...", H, H))

        Q, R = torch.qr(H)
        Q = torch.diag_embed(torch.sign(torch.diagonal(R, dim1=1, dim2=2))) @ Q

        # Normalize if need be. Uniform NormDistribution does nothing, Q is already orthonormal
        if norm_distribution == NormDistribution.Xi:
            return torch.diag_embed(norms) @ Q

        return Q
//  [python]
def _get_random_ortho_matrix(
        blocks: int,
        dim: int,
        device: torch.device,
        norm_distribution: NormDistribution = NormDistribution.Uniform,
    ) -> torch.Tensor:
        r"""
        Generate a random matrix whose rows are exactly orthonormal

        "How to generate random matrices from the classical compact groups", Mezzadri, 2007
        https://arxiv.org/pdf/math-ph/0609050v2.pdf

        .. note: the typical qr decomposition does not give uniform results, qr decomposition is not
        unique and the qr decomposition routines are biased towards numerical stability. See the above
        paper for more information.

        .. note: this does not follow the original implementation from the Performers authors.
        see docs/assets/kde plots to visualize the impact of using the R signs to correct Q
        """

        H = torch.randn((blocks, dim, dim), device=device, requires_grad=False)

        # Randomly scale the norms of the features, Xi distributed
        if norm_distribution == NormDistribution.Xi:
            # NOTE: This averages to sqrt(d)
            norms = torch.sqrt(torch.einsum("...d,...d->...", H, H))

        Q, R = torch.linalg.qr(H)
        Q = torch.diag_embed(torch.sign(torch.diagonal(R, dim1=1, dim2=2))) @ Q

        # Normalize if need be. Uniform NormDistribution does nothing, Q is already orthonormal
        if norm_distribution == NormDistribution.Xi:
            return torch.diag_embed(norms) @ Q

        return Q
//  [python]
def random_ortho(n: int, m: int) -> Tensor:
    """Generate a random orthogonal matrix.

    Args:
        n: Number of rows in the output matrix
        m: Number of columns in the output matrix

    Returns:
        A random orthogonal matrix of size (n, m)
    """
    q, _ = torch.qr(torch.randn(n, m))
    return q
//  [python]
def random_ortho(n: int, m: int) -> Tensor:
    """Generate a random orthogonal matrix.

    Args:
        n: Number of rows in the output matrix
        m: Number of columns in the output matrix

    Returns:
        A random orthogonal matrix of size (n, m)
    """
    q, _ = torch.linalg.qr(torch.randn(n, m))
    return q
//  [python]
def rand(self, num=1):
        if self.n == 2:
            sphere_point = torch.randn((num, 4), dtype=torch.double, device=device)
            sphere_point /= torch.linalg.vector_norm(sphere_point, dim=-1, keepdim=True)
            a = torch.view_as_complex(sphere_point[..., :2]).unsqueeze(-1)
            b = torch.view_as_complex(sphere_point[..., 2:]).unsqueeze(-1)
            r1 = torch.hstack((a, -b.conj())).unsqueeze(-1)
            r2 = torch.hstack((b, a.conj())).unsqueeze(-1)
            q = torch.cat((r1, r2), -1)
            return q
        else:
            h = torch.randn((num, self.n, self.n), dtype=dtype, device=device)
            q, r = torch.qr(h)
            r_diag = torch.diagonal(r, dim1=-2, dim2=-1)
            r_diag_inv_phase = torch.conj(r_diag / torch.abs(r_diag))
            q *= r_diag_inv_phase[:, None]
            q_det = torch.det(q)
            q_det_inv_phase = torch.conj(q_det / torch.abs(q_det))
            q[:, :, 0] *= q_det_inv_phase[:, None]
            return q
//  [python]
def rand(self, num=1):
        if self.n == 2:
            sphere_point = torch.randn((num, 4), dtype=torch.double, device=device)
            sphere_point /= torch.linalg.vector_norm(sphere_point, dim=-1, keepdim=True)
            a = torch.view_as_complex(sphere_point[..., :2]).unsqueeze(-1)
            b = torch.view_as_complex(sphere_point[..., 2:]).unsqueeze(-1)
            r1 = torch.hstack((a, -b.conj())).unsqueeze(-1)
            r2 = torch.hstack((b, a.conj())).unsqueeze(-1)
            q = torch.cat((r1, r2), -1)
            return q
        else:
            h = torch.randn((num, self.n, self.n), dtype=dtype, device=device)
            q, r = torch.linalg.qr(h)
            r_diag = torch.diagonal(r, dim1=-2, dim2=-1)
            r_diag_inv_phase = torch.conj(r_diag / torch.abs(r_diag))
            q *= r_diag_inv_phase[:, None]
            q_det = torch.det(q)
            q_det_inv_phase = torch.conj(q_det / torch.abs(q_det))
            q[:, :, 0] *= q_det_inv_phase[:, None]
            return q
//  [python]
def qr(x, mode="reduced"):
    if mode not in {"reduced", "complete"}:
        raise ValueError(
            "`mode` argument value not supported. "
            "Expected one of {'reduced', 'complete'}. "
            f"Received: mode={mode}"
        )
    return torch.qr(x, mode=mode)
//  [python]
def qr(x, mode="reduced"):
    if mode not in {"reduced", "complete"}:
        raise ValueError(
            "`mode` argument value not supported. "
            "Expected one of {'reduced', 'complete'}. "
            f"Received: mode={mode}"
        )
    return torch.linalg.qr(x, mode=mode)
//  [python]
def __init__(self, num_channels, use_lu=True):
        """Constructor

        Args:
          num_channels: Number of channels of the data
          use_lu: Flag whether to parametrize weights through the LU decomposition
        """
        super().__init__()
        self.num_channels = num_channels
        self.use_lu = use_lu
        Q, _ = torch.qr(torch.randn(self.num_channels, self.num_channels))
        if use_lu:
            P, L, U = torch.lu_unpack(*Q.lu())
            self.register_buffer("P", P)  # remains fixed during optimization
            self.L = nn.Parameter(L)  # lower triangular portion
            S = U.diag()  # "crop out" the diagonal to its own parameter
            self.register_buffer("sign_S", torch.sign(S))
            self.log_S = nn.Parameter(torch.log(torch.abs(S)))
            self.U = nn.Parameter(
                torch.triu(U, diagonal=1)
            )  # "crop out" diagonal, stored in S
            self.register_buffer("eye", torch.diag(torch.ones(self.num_channels)))
        else:
            self.W = nn.Parameter(Q)
//  [python]
def __init__(self, num_channels, use_lu=True):
        """Constructor

        Args:
          num_channels: Number of channels of the data
          use_lu: Flag whether to parametrize weights through the LU decomposition
        """
        super().__init__()
        self.num_channels = num_channels
        self.use_lu = use_lu
        Q, _ = torch.linalg.qr(torch.randn(self.num_channels, self.num_channels))
        if use_lu:
            P, L, U = torch.lu_unpack(*Q.lu())
            self.register_buffer("P", P)  # remains fixed during optimization
            self.L = nn.Parameter(L)  # lower triangular portion
            S = U.diag()  # "crop out" the diagonal to its own parameter
            self.register_buffer("sign_S", torch.sign(S))
            self.log_S = nn.Parameter(torch.log(torch.abs(S)))
            self.U = nn.Parameter(
                torch.triu(U, diagonal=1)
            )  # "crop out" diagonal, stored in S
            self.register_buffer("eye", torch.diag(torch.ones(self.num_channels)))
        else:
            self.W = nn.Parameter(Q)
//  [python]
def __init__(self, num_channels, use_lu=False):
        """Constructor

        Args:
          num_channels: Number of channels of the data
          use_lu: Flag whether to parametrize weights through the LU decomposition
        """
        super().__init__()
        self.num_channels = num_channels
        self.use_lu = use_lu
        Q, _ = torch.qr(torch.randn(self.num_channels, self.num_channels))
        if use_lu:
            P, L, U = torch.lu_unpack(*Q.lu())
            self.register_buffer("P", P)  # remains fixed during optimization
            self.L = nn.Parameter(L)  # lower triangular portion
            S = U.diag()  # "crop out" the diagonal to its own parameter
            self.register_buffer("sign_S", torch.sign(S))
            self.log_S = nn.Parameter(torch.log(torch.abs(S)))
            self.U = nn.Parameter(
                torch.triu(U, diagonal=1)
            )  # "crop out" diagonal, stored in S
            self.register_buffer("eye", torch.diag(torch.ones(self.num_channels)))
        else:
            self.W = nn.Parameter(Q)
//  [python]
def __init__(self, num_channels, use_lu=False):
        """Constructor

        Args:
          num_channels: Number of channels of the data
          use_lu: Flag whether to parametrize weights through the LU decomposition
        """
        super().__init__()
        self.num_channels = num_channels
        self.use_lu = use_lu
        Q, _ = torch.linalg.qr(torch.randn(self.num_channels, self.num_channels))
        if use_lu:
            P, L, U = torch.lu_unpack(*Q.lu())
            self.register_buffer("P", P)  # remains fixed during optimization
            self.L = nn.Parameter(L)  # lower triangular portion
            S = U.diag()  # "crop out" the diagonal to its own parameter
            self.register_buffer("sign_S", torch.sign(S))
            self.log_S = nn.Parameter(torch.log(torch.abs(S)))
            self.U = nn.Parameter(
                torch.triu(U, diagonal=1)
            )  # "crop out" diagonal, stored in S
            self.register_buffer("eye", torch.diag(torch.ones(self.num_channels)))
        else:
            self.W = nn.Parameter(Q)
//  [python]
def interpolate_eigvecs(U_1: torch.Tensor, U_2: torch.Tensor, mask: torch.Tensor, alpha: torch.Tensor):
    """ Canonical Ordering of Eigenvecotrs (direction is flipped such that max absolute value is made positive and node order is picked by sorting by eigenvectors in order)
    U_1 and U_2: [batch_size (N), num_objects (L), num_eigvecs(C)] - assumed to be sorted with sort_eigvecs
    mask: [batch_size (N), num_objects (L), 1]
    alpha: [batch_size (N), num_objects (L), 1] - Value of 1.0 returns U_1, while value of 0.0 retursn U_2
    """ 
    U = U_2 + ((U_1 - U_2) * alpha) # if alpha=1 take first matrix (U_1)
    U, _ = torch.linalg.qr(U, mode='reduced')
    # U = gram_schmidt(U)
    U = deterministic_vector_sign_flip(U)
    U = U * mask
    return U
//  [python]
def interpolate_eigvecs(U_1: torch.Tensor, U_2: torch.Tensor, mask: torch.Tensor, alpha: torch.Tensor):
    """ Canonical Ordering of Eigenvecotrs (direction is flipped such that max absolute value is made positive and node order is picked by sorting by eigenvectors in order)
    U_1 and U_2: [batch_size (N), num_objects (L), num_eigvecs(C)] - assumed to be sorted with sort_eigvecs
    mask: [batch_size (N), num_objects (L), 1]
    alpha: [batch_size (N), num_objects (L), 1] - Value of 1.0 returns U_1, while value of 0.0 retursn U_2
    """ 
    U = U_2 + ((U_1 - U_2) * alpha) # if alpha=1 take first matrix (U_1)
    U = torch.linalg.qr(U, mode='reduced')[0]
    # U = gram_schmidt(U)
    U = deterministic_vector_sign_flip(U)
    U = U * mask
    return U
//  [python]
def ql_decomposition(A):
    P = torch.tensor([[0, 0, 1], [0, 1, 0], [1, 0, 0]], device=A.device).float()
    A_tilde = torch.matmul(A, P)
    Q_tilde, R_tilde = torch.qr(A_tilde)
    Q = torch.matmul(Q_tilde, P)
    L = torch.matmul(torch.matmul(P, R_tilde), P)
    d = torch.diag(L)
    Q[:, 0] *= torch.sign(d[0])
    Q[:, 1] *= torch.sign(d[1])
    Q[:, 2] *= torch.sign(d[2])
    L[0] *= torch.sign(d[0])
    L[1] *= torch.sign(d[1])
    L[2] *= torch.sign(d[2])
    return Q, L
//  [python]
def ql_decomposition(A):
    P = torch.tensor([[0, 0, 1], [0, 1, 0], [1, 0, 0]], device=A.device).float()
    A_tilde = torch.matmul(A, P)
    Q_tilde, R_tilde = torch.linalg.qr(A_tilde)
    Q = torch.matmul(Q_tilde, P)
    L = torch.matmul(torch.matmul(P, R_tilde), P)
    d = torch.diag(L)
    Q[:, 0] *= torch.sign(d[0])
    Q[:, 1] *= torch.sign(d[1])
    Q[:, 2] *= torch.sign(d[2])
    L[0] *= torch.sign(d[0])
    L[1] *= torch.sign(d[1])
    L[2] *= torch.sign(d[2])
    return Q, L
//  [python]
def qr(matrix, full_matrices=False):
    mode = "reduced" if not full_matrices else "complete"
    Q, R = torch.qr(matrix, mode=mode)
    return Q, R
//  [python]
def qr(matrix, full_matrices=False):
    mode = "reduced" if not full_matrices else "complete"
    Q, R = torch.linalg.qr(matrix, mode=mode)
    return Q, R
//  [python]
def __init__(self, dim=2):
        super().__init__()

        w = torch.randn(dim, dim)
        w, _ = torch.qr(w)   # W^{-1} = W^T (only at initialization)
        self.w = nn.Parameter(w)
//  [python]
def __init__(self, dim=2):
        super().__init__()

        w = torch.randn(dim, dim)
        w = torch.linalg.qr(w)[0]   # W^{-1} = W^T (only at initialization)
        self.w = nn.Parameter(w)
//  [python]
def to_stiefel_qr(theta, dim:int, rank:int):
    r'''map real vector to a Stiefel manifold via QR decomposition

    Parameters:
        theta (np.ndarray,torch.Tensor): the last dimension will be expanded to the matrix
                and the rest dimensions will be batch dimensions. For real case, the last dimension
                should be `dim*rank`, and for complex case, the last dimension should be `2*dim*rank`.
        dim (int): dimension of the matrix.
        rank (int): rank of the matrix.

    Returns:
        ret (np.ndarray,torch.Tensor): array of shape `theta.shape[:-1]+(dim,rank)`
    '''
    # TODO poor efficiency in L-BFGS-B optimization
    # the forward and backward time is still okay compared with matrix-square-root,
    # but the time for L-BFGS-B optimization is much longer (strange)
    assert dim>=rank
    is_torch = isinstance(theta, torch.Tensor)
    shape = theta.shape
    if shape[-1]==dim*rank: #real
        mat = theta.reshape(-1,dim,rank)
    else: #complex
        assert shape[-1]==2*dim*rank
        tmp0 = theta.reshape(-1,2,dim,rank)
        if is_torch:
            mat = torch.complex(tmp0[:,0], tmp0[:,1])
        else:
            mat = tmp0[:,0] + 1j* tmp0[:,1]
    mat = mat.reshape(*shape[:-1], dim, rank)
    # TODO poor efficiency in L-BFGS-B optimization
    # the forward and backward time is still okay compared with matrix-cholesky,
    # but the time for L-BFGS-B optimization is much longer (strange)
    if isinstance(mat, torch.Tensor):
        q, r = torch.qr(mat)
        ret = q
    else:
        ret = np.linalg.qr(mat, mode='reduced')[0]
    return ret
//  [python]
def to_stiefel_qr(theta, dim:int, rank:int):
    r'''map real vector to a Stiefel manifold via QR decomposition

    Parameters:
        theta (np.ndarray,torch.Tensor): the last dimension will be expanded to the matrix
                and the rest dimensions will be batch dimensions. For real case, the last dimension
                should be `dim*rank`, and for complex case, the last dimension should be `2*dim*rank`.
        dim (int): dimension of the matrix.
        rank (int): rank of the matrix.

    Returns:
        ret (np.ndarray,torch.Tensor): array of shape `theta.shape[:-1]+(dim,rank)`
    '''
    # TODO poor efficiency in L-BFGS-B optimization
    # the forward and backward time is still okay compared with matrix-square-root,
    # but the time for L-BFGS-B optimization is much longer (strange)
    assert dim>=rank
    is_torch = isinstance(theta, torch.Tensor)
    shape = theta.shape
    if shape[-1]==dim*rank: #real
        mat = theta.reshape(-1,dim,rank)
    else: #complex
        assert shape[-1]==2*dim*rank
        tmp0 = theta.reshape(-1,2,dim,rank)
        if is_torch:
            mat = torch.complex(tmp0[:,0], tmp0[:,1])
        else:
            mat = tmp0[:,0] + 1j* tmp0[:,1]
    mat = mat.reshape(*shape[:-1], dim, rank)
    # TODO poor efficiency in L-BFGS-B optimization
    # the forward and backward time is still okay compared with matrix-cholesky,
    # but the time for L-BFGS-B optimization is much longer (strange)
    if isinstance(mat, torch.Tensor):
        ret = torch.linalg.qr(mat, mode='reduced')[0]
    else:
        ret = np.linalg.qr(mat, mode='reduced')[0]
    return ret
//  [python]
def qr_retraction(tan_vec):  # tan_vec, p-by-n, p <= n
    [p, n] = tan_vec.size()
    tan_vec.t_()
    q, r = torch.qr(tan_vec)
    d = torch.diag(r, 0)
    ph = d.sign()
    q *= ph.expand_as(q)
    q.t_()

    return q
//  [python]
def qr_retraction(tan_vec):  # tan_vec, p-by-n, p <= n
    [p, n] = tan_vec.size()
    tan_vec.t_()
    q, r = torch.linalg.qr(tan_vec)
    d = torch.diag(r, 0)
    ph = d.sign()
    q *= ph.expand_as(q)
    q.t_()

    return q
//  [python]
def _init_cache_for_non_constant_diag(self, eye: Tensor, batch_shape: Union[torch.Size, List[int]], n: int):
        # With non-constant diagonals, we cant factor out the noise as easily
        self._q_cache, self._r_cache = torch.qr(
            torch.cat((self._piv_chol_self / self._noise.sqrt(), eye), dim=-2)
        )
        self._q_cache = self._q_cache[..., :n, :] / self._noise.sqrt()

        # Use the matrix determinant lemma for the logdet, using the fact that R'R = L_k'L_k + s*I
        logdet = self._r_cache.diagonal(dim1=-1, dim2=-2).abs().log().sum(-1).mul(2)
        logdet -= (1.0 / self._noise).log().sum([-1, -2])
        self._precond_logdet_cache = logdet.view(*batch_shape) if len(batch_shape) else logdet.squeeze()
//  [python]
def _init_cache_for_non_constant_diag(self, eye: Tensor, batch_shape: Union[torch.Size, List[int]], n: int):
        # With non-constant diagonals, we cant factor out the noise as easily
        self._q_cache, self._r_cache = torch.linalg.qr(
            torch.cat((self._piv_chol_self / self._noise.sqrt(), eye), dim=-2)
        )
        self._q_cache = self._q_cache[..., :n, :] / self._noise.sqrt()

        # Use the matrix determinant lemma for the logdet, using the fact that R'R = L_k'L_k + s*I
        logdet = self._r_cache.diagonal(dim1=-1, dim2=-2).abs().log().sum(-1).mul(2)
        logdet -= (1.0 / self._noise).log().sum([-1, -2])
        self._precond_logdet_cache = logdet.view(*batch_shape) if len(batch_shape) else logdet.squeeze()
//  [python]
def _init_cache_for_constant_diag(self, eye: Tensor, batch_shape: Union[torch.Size, List[int]], n: int, k: int):
        # We can factor out the noise for for both QR and solves.
        self._noise = self._noise.narrow(-2, 0, 1)
        self._q_cache, self._r_cache = torch.qr(
            torch.cat((self._piv_chol_self, self._noise.sqrt() * eye), dim=-2)
        )
        self._q_cache = self._q_cache[..., :n, :]

        # Use the matrix determinant lemma for the logdet, using the fact that R'R = L_k'L_k + s*I
        logdet = self._r_cache.diagonal(dim1=-1, dim2=-2).abs().log().sum(-1).mul(2)
        logdet = logdet + (n - k) * self._noise.squeeze(-2).squeeze(-1).log()
        self._precond_logdet_cache = logdet.view(*batch_shape) if len(batch_shape) else logdet.squeeze()
//  [python]
def _init_cache_for_constant_diag(self, eye: Tensor, batch_shape: Union[torch.Size, List[int]], n: int, k: int):
        # We can factor out the noise for for both QR and solves.
        self._noise = self._noise.narrow(-2, 0, 1)
        self._q_cache, self._r_cache = torch.linalg.qr(
            torch.cat((self._piv_chol_self, self._noise.sqrt() * eye), dim=-2)
        )
        self._q_cache = self._q_cache[..., :n, :]

        # Use the matrix determinant lemma for the logdet, using the fact that R'R = L_k'L_k + s*I
        logdet = self._r_cache.diagonal(dim1=-1, dim2=-2).abs().log().sum(-1).mul(2)
        logdet = logdet + (n - k) * self._noise.squeeze(-2).squeeze(-1).log()
        self._precond_logdet_cache = logdet.view(*batch_shape) if len(batch_shape) else logdet.squeeze()