jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def round_ldl_gptqequiv(
    w,
    H,
    nbits,
    unbiased=False
):
    """
    w in R^{m,d}
    d: input_shape, m: output_shape
    """
    (d, d_) = H.shape
    assert (d == d_)
    (m, d) = w.shape
    H = torch.flip(H, [0,1])
    L = torch.cholesky(H)
    L = torch.flip(L,[0,1])
    L = L @ torch.diag(1/torch.diag(L))
    L = L - torch.eye(d, device=L.device)
    if unbiased:
        eta = torch.rand(w.shape).to(w.device)
    else:
        eta = 0.5 * torch.ones(w.shape).to(w.device)
    w_hat = w.clone()
    for i in range(d):
        #w_hat[:,i] = torch.clamp(torch.floor(w[:,i] + (w - w_hat) @ L[:,i] + eta[:,i]), min=0, max=2**nbits-1)
        # optimized version
        w_hat[:,i] = torch.clamp(torch.floor(w[:,i] + (w[:,:i+1] - w_hat[:,:i+1]) @ L[:i+1,i] + eta[:,i]), min=0, max=2**nbits-1)

    # reverse order
    # L = torch.cholesky(H)
    # L = L @ torch.diag(1/torch.diag(L))
    # if unbiased:
    #     eta = torch.rand(w.shape).to(w.device)
    # else:
    #     eta = 0.5 * torch.ones(w.shape).to(w.device)
    # w_hat = w.clone()
    # for i in reversed(range(d)):
    #     w_hat[:,i] = torch.clamp(torch.floor(w[:,i] + (w[:,i:] - w_hat[:,i:]) @ L[i:,i] + eta[:,i]), min=0, max=2**nbits-1)

    wr = w_hat
    wr_counts = check_nbits(wr, nbits)
    return wr
//  [python]
def round_ldl_gptqequiv(
    w,
    H,
    nbits,
    unbiased=False
):
    """
    w in R^{m,d}
    d: input_shape, m: output_shape
    """
    (d, d_) = H.shape
    assert (d == d_)
    (m, d) = w.shape
    H = torch.flip(H, [0,1])
    L = torch.linalg.cholesky(H)
    L = torch.flip(L,[0,1])
    L = L @ torch.diag(1/torch.diag(L))
    L = L - torch.eye(d, device=L.device)
    if unbiased:
        eta = torch.rand(w.shape).to(w.device)
    else:
        eta = 0.5 * torch.ones(w.shape).to(w.device)
    w_hat = w.clone()
    for i in range(d):
        #w_hat[:,i] = torch.clamp(torch.floor(w[:,i] + (w - w_hat) @ L[:,i] + eta[:,i]), min=0, max=2**nbits-1)
        # optimized version
        w_hat[:,i] = torch.clamp(torch.floor(w[:,i] + (w[:,:i+1] - w_hat[:,:i+1]) @ L[:i+1,i] + eta[:,i]), min=0, max=2**nbits-1)

    # reverse order
    # L = torch.linalg.cholesky(H)
    # L = L @ torch.diag(1/torch.diag(L))
    # if unbiased:
    #     eta = torch.rand(w.shape).to(w.device)
    # else:
    #     eta = 0.5 * torch.ones(w.shape).to(w.device)
    # w_hat = w.clone()
    # for i in reversed(range(d)):
    #     w_hat[:,i] = torch.clamp(torch.floor(w[:,i] + (w[:,i:] - w_hat[:,i:]) @ L[i:,i] + eta[:,i]), min=0, max=2**nbits-1)

    wr = w_hat
    wr_counts = check_nbits(wr, nbits)
    return wr
//  [python]
def test_cholesky_inverse(input_data):
    m = torch.cholesky_inverse
    inp = torch.mm(input_data, input_data.t()) + 1e-05 * torch.eye(
        input_data.shape[0]
    )  # make symmetric positive definite
    u = inp.cholesky()  # pylint: disable=C0103
    output = m(u)
    output_musa = m(u.musa())
    assert testing.DefaultComparator(abs_diff=1e-5)(output, output_musa)
//  [python]
def test_cholesky_inverse(input_data):
    m = torch.cholesky_inverse
    inp = torch.mm(input_data, input_data.t()) + 1e-05 * torch.eye(
        input_data.shape[0]
    )  # make symmetric positive definite
    u = torch.linalg.cholesky(inp)  # pylint: disable=C0103
    output = m(u)
    output_musa = m(u.musa())
    assert testing.DefaultComparator(abs_diff=1e-5)(output, output_musa)
//  [python]
def log_likelihood(self, x_train, y_train):
        K = self.kernel(x_train, x_train) + self.noise_variance.pow(2) * torch.eye(len(x_train))
        
        L = torch.cholesky(K)
        log_det_K = 2 * torch.sum(torch.log(torch.diag(L)))
        Alpha = torch.cholesky_solve(y_train, L, upper = False)
        
        # return - 0.5 * (Alpha.T @ Alpha + log_det_K + len(x_train) * np.log(2 * np.pi))
        return - 0.5 * ( (Alpha ** 2).sum() + log_det_K + len(x_train) * np.log(2 * np.pi))
//  [python]
def log_likelihood(self, x_train, y_train):
        K = self.kernel(x_train, x_train) + self.noise_variance.pow(2) * torch.eye(len(x_train))
        
        L = torch.linalg.cholesky(K)
        log_det_K = 2 * torch.sum(torch.log(torch.diag(L)))
        Alpha = torch.cholesky_solve(y_train, L, upper = False)
        
        # return - 0.5 * (Alpha.T @ Alpha + log_det_K + len(x_train) * np.log(2 * np.pi))
        return - 0.5 * ( (Alpha ** 2).sum() + log_det_K + len(x_train) * np.log(2 * np.pi))
//  [python]
def round_ldl_block(
    w,
    H,
    nbits,
    blocksize=128,
    n_greedy_passes=9,
    unbiased=False
):
    """
    w in R^{m,d}
    d: input_shape, m: output_shape
    note: this has been updated with a (hopefully) more efficient LDL pass
    """
    assert (not unbiased) or (n_greedy_passes == 0), "greedy passes are incompatible with unbiased LDL rounding"
    (d, d_) = H.shape
    assert (d == d_)
    (m, d) = w.shape
    L = torch.cholesky(H)
    L = L @ torch.diag(1/torch.diag(L))
    L = L - torch.eye(d, device=L.device)
    if unbiased:
        eta = torch.rand(w.shape).to(w.device)
    else:
        eta = 0.5 * torch.ones(w.shape).to(w.device)
    w_hat = w.clone()
    for i2 in range(d, 0, -blocksize):
        i1 = max(i2 - blocksize, 0)
        count = i2 - i1
        W1 = w[:, i1:i2]
        W2Hdiff = w[:, i2:] - w_hat[:, i2:]
        WHat1 = w_hat[:, i1:i2].clone()
        L1 = L[:, i1:i2]
        Eta1 = eta[:, i1:i2]

        for i in reversed(range(count)):
            WHat1[:,i] = torch.clamp(
                torch.floor(W1[:,i] + (W1 - WHat1) @ L1[i1:i2,i] + W2Hdiff @ L1[i2:,i] + Eta1[:,i]), 
                min=0, max=2**nbits-1)

        w_hat[:, i1:i2] = WHat1

    wr = w_hat.clone()
    s = w_hat - w
    H = H / H.diag().max()

    for igp in range(n_greedy_passes):
        for i2 in range(d, 0, -blocksize):
            i1 = max(i2 - blocksize, 0)
            count = i2 - i1
            W1 = wr[:, i1:i2].clone()
            S0 = s[:, :i1]
            S1 = s[:, i1:i2].clone()
            S2 = s[:, i2:]
            H0 = H[:i1, i1:i2]
            H1 = H[i1:i2, i1:i2]
            H2 = H[i2:, i1:i2]

            for i in reversed(range(count)):
                Hs = S0 @ H0[:, i] + S1 @ H1[:, i] + S2 @ H2[:, i]
                epsXTsj = W1[:, i] - torch.round(W1[:, i] - Hs / H1[i,i])
                W1[:, i] -= epsXTsj
                S1[:, i] -= epsXTsj

            wr[:, i1:i2] = W1
            s[:, i1:i2] = S1

        wr = torch.clamp(wr, min=0, max=2**nbits - 1)
        if ((w_hat == wr).all()):
            sys.stderr.write(f"breaking after {igp+1} greedy passes found fixed point")
            break
        w_hat.copy_(wr)
    
    wr_counts = check_nbits(wr, nbits)
    return wr
//  [python]
def round_ldl_block(
    w,
    H,
    nbits,
    blocksize=128,
    n_greedy_passes=9,
    unbiased=False
):
    """
    w in R^{m,d}
    d: input_shape, m: output_shape
    note: this has been updated with a (hopefully) more efficient LDL pass
    """
    assert (not unbiased) or (n_greedy_passes == 0), "greedy passes are incompatible with unbiased LDL rounding"
    (d, d_) = H.shape
    assert (d == d_)
    (m, d) = w.shape
    L = torch.linalg.cholesky(H)
    L = L @ torch.diag(1/torch.diag(L))
    L = L - torch.eye(d, device=L.device)
    if unbiased:
        eta = torch.rand(w.shape).to(w.device)
    else:
        eta = 0.5 * torch.ones(w.shape).to(w.device)
    w_hat = w.clone()
    for i2 in range(d, 0, -blocksize):
        i1 = max(i2 - blocksize, 0)
        count = i2 - i1
        W1 = w[:, i1:i2]
        W2Hdiff = w[:, i2:] - w_hat[:, i2:]
        WHat1 = w_hat[:, i1:i2].clone()
        L1 = L[:, i1:i2]
        Eta1 = eta[:, i1:i2]

        for i in reversed(range(count)):
            WHat1[:,i] = torch.clamp(
                torch.floor(W1[:,i] + (W1 - WHat1) @ L1[i1:i2,i] + W2Hdiff @ L1[i2:,i] + Eta1[:,i]), 
                min=0, max=2**nbits-1)

        w_hat[:, i1:i2] = WHat1

    wr = w_hat.clone()
    s = w_hat - w
    H = H / H.diag().max()

    for igp in range(n_greedy_passes):
        for i2 in range(d, 0, -blocksize):
            i1 = max(i2 - blocksize, 0)
            count = i2 - i1
            W1 = wr[:, i1:i2].clone()
            S0 = s[:, :i1]
            S1 = s[:, i1:i2].clone()
            S2 = s[:, i2:]
            H0 = H[:i1, i1:i2]
            H1 = H[i1:i2, i1:i2]
            H2 = H[i2:, i1:i2]

            for i in reversed(range(count)):
                Hs = S0 @ H0[:, i] + S1 @ H1[:, i] + S2 @ H2[:, i]
                epsXTsj = W1[:, i] - torch.round(W1[:, i] - Hs / H1[i,i])
                W1[:, i] -= epsXTsj
                S1[:, i] -= epsXTsj

            wr[:, i1:i2] = W1
            s[:, i1:i2] = S1

        wr = torch.clamp(wr, min=0, max=2**nbits - 1)
        if ((w_hat == wr).all()):
            sys.stderr.write(f"breaking after {igp+1} greedy passes found fixed point")
            break
        w_hat.copy_(wr)
    
    wr_counts = check_nbits(wr, nbits)
    return wr
//  [python]
def __call__(self, data):
        if not self.patch_size:
            raise ValueError("Patch size must be defined for noise prewhitening.")

        if data.shape[-1] != 2:
            data = torch.view_as_real(data)

        noise = data[:, self.patch_size[0] : self.patch_size[1], self.patch_size[-2] : self.patch_size[-1]]
        noise_int = torch.reshape(noise, (noise.shape[0], int(torch.numel(noise) / noise.shape[0])))

        deformation_matrix = (1 / (float(noise_int.shape[1]) - 1)) * torch.mm(noise_int, torch.conj(noise_int).t())
        psi = torch.inverse(torch.cholesky(deformation_matrix)) * sqrt(2) * sqrt(self.scale_factor)

        return torch.reshape(
            torch.mm(psi, torch.reshape(data, (data.shape[0], int(torch.numel(data) / data.shape[0])))), data.shape
        )
//  [python]
def __call__(self, data):
        if not self.patch_size:
            raise ValueError("Patch size must be defined for noise prewhitening.")

        if data.shape[-1] != 2:
            data = torch.view_as_real(data)

        noise = data[:, self.patch_size[0] : self.patch_size[1], self.patch_size[-2] : self.patch_size[-1]]
        noise_int = torch.reshape(noise, (noise.shape[0], int(torch.numel(noise) / noise.shape[0])))

        deformation_matrix = (1 / (float(noise_int.shape[1]) - 1)) * torch.mm(noise_int, torch.conj(noise_int).t())
        psi = torch.linalg.inv(torch.linalg.cholesky(deformation_matrix)) * sqrt(2) * sqrt(self.scale_factor)

        return torch.reshape(
            torch.mm(psi, torch.reshape(data, (data.shape[0], int(torch.numel(data) / data.shape[0])))), data.shape
        )
//  [python]
def cholesky_inv(X, damping=1e-7):
    diag = torch.diagonal(X)
    diag += damping
    u = torch.cholesky(X)
    diag -= damping
    return torch.cholesky_inverse(u)
//  [python]
def cholesky_inv(X, damping=1e-7):
    diag = torch.diagonal(X)
    diag += damping
    u = torch.linalg.cholesky(X)
    diag -= damping
    return torch.cholesky_inverse(u)
//  [python]
def set(self, distr):   
        """Sets distribution"""
        assert 'mean' in distr and 'cov' in distr
        
        self.distr = distr.copy()
        self.mean = distr['mean']
        assert isinstance(self.mean, np.ndarray)

        self.ndims = len(self.mean)
        self.cov = distr['cov']
        assert isinstance(self.cov, np.ndarray)

        # Bounds for denormalization
        self.low, self.high = None, None
        if 'low' in distr:
            assert 'high' in distr
            # Assume samples to be normalized from [low, high] to [0 , 1]
            # Save bounds to potentially denormalize samples
            self.low = np.array(distr['low'])
            self.high = np.array(distr['high'])

        # Transform sigma into L
        self.to_mean = torch.tensor(self.mean)

        to_cov = torch.tensor(self.cov)
        L = torch.cholesky(to_cov).double()
        D = L.size(0)
        mask = torch.ones(D, D).tril(-1).bool()  # all elements strictly below the main diagonal
        self.L_low_diag = torch.zeros(D, D).double()

        self.L_off_diag = torch.tensor(L[mask])  # off-diagonal elements of the lower-side
        self.log_diag_L = torch.log(L.diag())  # log-transformed elements of diagional

        self.to_mean.requires_grad_(True)     # optimize means
        self.L_off_diag.requires_grad_(True)  # optimize lower off-diagonal elements of L
        self.log_diag_L.requires_grad_(True)  # optimize log-diagonal elements of L

        self.L_low_diag[mask] = self.L_off_diag
        self.to_L = self.L_low_diag + torch.eye(D)*torch.exp(self.log_diag_L)  # propagates gradients correctly, instead of using L
        self.L = self.to_L.clone().detach()

        self.to_distr = MultivariateNormal(loc=self.to_mean,
                                           covariance_matrix=None,
                                           precision_matrix=None,
                                           scale_tril=self.to_L)


        self.to_params = [self.to_mean, self.L_off_diag, self.log_diag_L]
//  [python]
def set(self, distr):   
        """Sets distribution"""
        assert 'mean' in distr and 'cov' in distr
        
        self.distr = distr.copy()
        self.mean = distr['mean']
        assert isinstance(self.mean, np.ndarray)

        self.ndims = len(self.mean)
        self.cov = distr['cov']
        assert isinstance(self.cov, np.ndarray)

        # Bounds for denormalization
        self.low, self.high = None, None
        if 'low' in distr:
            assert 'high' in distr
            # Assume samples to be normalized from [low, high] to [0 , 1]
            # Save bounds to potentially denormalize samples
            self.low = np.array(distr['low'])
            self.high = np.array(distr['high'])

        # Transform sigma into L
        self.to_mean = torch.tensor(self.mean)

        to_cov = torch.tensor(self.cov)
        L = torch.linalg.cholesky(to_cov).double()
        D = L.size(0)
        mask = torch.ones(D, D).tril(-1).bool()  # all elements strictly below the main diagonal
        self.L_low_diag = torch.zeros(D, D).double()

        self.L_off_diag = torch.tensor(L[mask])  # off-diagonal elements of the lower-side
        self.log_diag_L = torch.log(L.diag())  # log-transformed elements of diagional

        self.to_mean.requires_grad_(True)     # optimize means
        self.L_off_diag.requires_grad_(True)  # optimize lower off-diagonal elements of L
        self.log_diag_L.requires_grad_(True)  # optimize log-diagonal elements of L

        self.L_low_diag[mask] = self.L_off_diag
        self.to_L = self.L_low_diag + torch.eye(D)*torch.exp(self.log_diag_L)  # propagates gradients correctly, instead of using L
        self.L = self.to_L.clone().detach()

        self.to_distr = MultivariateNormal(loc=self.to_mean,
                                           covariance_matrix=None,
                                           precision_matrix=None,
                                           scale_tril=self.to_L)


        self.to_params = [self.to_mean, self.L_off_diag, self.log_diag_L]
//  [python]
def cholesky_solve(X, b, damping=1e-7):
    diag = torch.diagonal(X)
    diag += damping
    u = torch.cholesky(X)
    diag -= damping
    return torch.cholesky_solve(b, u)
//  [python]
def cholesky_solve(X, b, damping=1e-7):
    diag = torch.diagonal(X)
    diag += damping
    u = torch.linalg.cholesky(X)
    diag -= damping
    return torch.cholesky_solve(b, u)
//  [python]
def generate_data(num_frames, num_sensors):
    """
    Generate data from a damped NCV dynamics model
    """
    dt = TIME_STEP
    bias_scale = 4.0
    obs_noise = 1.0
    trans_noise = 0.3

    # define dynamics
    z = torch.cat([10.0 * torch.randn(2), torch.rand(2)])  # position  # velocity
    damp = 0.1  # damp the velocities
    f = torch.tensor(
        [
            [1, 0, 0, 0],
            [0, 1, 0, 0],
            [dt * math.exp(-damp * dt), 0, math.exp(-damp * dt), 0],
            [0, dt * math.exp(-damp * dt), 0, math.exp(-damp * dt)],
        ]
    )
    trans_dist = dist.MultivariateNormal(
        torch.zeros(4),
        scale_tril=trans_noise * torch.cholesky(NCV_PROCESS_NOISE),
    )

    # define biased sensors
    sensor_bias = bias_scale * torch.randn(2 * num_sensors)
    h = torch.eye(4, 2).unsqueeze(-1).expand(-1, -1, num_sensors).reshape(4, -1)
    obs_dist = dist.MultivariateNormal(
        sensor_bias, scale_tril=obs_noise * torch.eye(2 * num_sensors)
    )

    states = []
    observations = []
    for t in range(num_frames):
        z = z @ f + trans_dist.sample()
        states.append(z)

        x = z @ h + obs_dist.sample()
        observations.append(x)

    states = torch.stack(states)
    observations = torch.stack(observations)
    assert observations.shape == (num_frames, num_sensors * 2)
    return observations, states, sensor_bias
//  [python]
def generate_data(num_frames, num_sensors):
    """
    Generate data from a damped NCV dynamics model
    """
    dt = TIME_STEP
    bias_scale = 4.0
    obs_noise = 1.0
    trans_noise = 0.3

    # define dynamics
    z = torch.cat([10.0 * torch.randn(2), torch.rand(2)])  # position  # velocity
    damp = 0.1  # damp the velocities
    f = torch.tensor(
        [
            [1, 0, 0, 0],
            [0, 1, 0, 0],
            [dt * math.exp(-damp * dt), 0, math.exp(-damp * dt), 0],
            [0, dt * math.exp(-damp * dt), 0, math.exp(-damp * dt)],
        ]
    )
    trans_dist = dist.MultivariateNormal(
        torch.zeros(4),
        scale_tril=trans_noise * torch.linalg.cholesky(NCV_PROCESS_NOISE),
    )

    # define biased sensors
    sensor_bias = bias_scale * torch.randn(2 * num_sensors)
    h = torch.eye(4, 2).unsqueeze(-1).expand(-1, -1, num_sensors).reshape(4, -1)
    obs_dist = dist.MultivariateNormal(
        sensor_bias, scale_tril=obs_noise * torch.eye(2 * num_sensors)
    )

    states = []
    observations = []
    for t in range(num_frames):
        z = z @ f + trans_dist.sample()
        states.append(z)

        x = z @ h + obs_dist.sample()
        observations.append(x)

    states = torch.stack(states)
    observations = torch.stack(observations)
    assert observations.shape == (num_frames, num_sensors * 2)
    return observations, states, sensor_bias
//  [python]
def fasterprune(
        self, sparsity, prune_n=0, prune_m=0, blocksize=128, percdamp=.01
    ):
        W = self.layer.weight.data.clone()
        if isinstance(self.layer, nn.Conv2d):
            W = W.flatten(1)
        if isinstance(self.layer, transformers.Conv1D):
            W = W.t()
        W = W.float()

        tick = time.time()

        H = self.H
        del self.H
        dead = torch.diag(H) == 0
        H[dead, dead] = 1
        W[:, dead] = 0

        Losses = torch.zeros(self.rows, device=self.dev)

        damp = percdamp * torch.mean(torch.diag(H))
        diag = torch.arange(self.columns, device=self.dev)
        H[diag, diag] += damp
        H = torch.cholesky(H)
        H = torch.cholesky_inverse(H)
        H = torch.cholesky(H, upper=True)
        Hinv = H

        mask = None

        for i1 in range(0, self.columns, blocksize):
            i2 = min(i1 + blocksize, self.columns)
            count = i2 - i1

            W1 = W[:, i1:i2].clone()
            Q1 = torch.zeros_like(W1)
            Err1 = torch.zeros_like(W1)
            Losses1 = torch.zeros_like(W1)
            Hinv1 = Hinv[i1:i2, i1:i2]

            if prune_n == 0: 
                if mask is not None:
                    mask1 = mask[:, i1:i2]
                else:
                    tmp = W1 ** 2 / (torch.diag(Hinv1).reshape((1, -1))) ** 2
                    thresh = torch.sort(tmp.flatten())[0][int(tmp.numel() * sparsity)]
                    mask1 = tmp <= thresh
            else:
                mask1 = torch.zeros_like(W1) == 1

            for i in range(count):
                w = W1[:, i]
                d = Hinv1[i, i]

                if prune_n != 0 and i % prune_m == 0:
                    tmp = W1[:, i:(i + prune_m)] ** 2 / (torch.diag(Hinv1)[i:(i + prune_m)].reshape((1, -1))) ** 2
                    mask1.scatter_(1, i + torch.topk(tmp, prune_n, dim=1, largest=False)[1], True)

                q = w.clone()
                q[mask1[:, i]] = 0

                Q1[:, i] = q
                Losses1[:, i] = (w - q) ** 2 / d ** 2

                err1 = (w - q) / d 
                W1[:, i:] -= err1.unsqueeze(1).matmul(Hinv1[i, i:].unsqueeze(0))
                Err1[:, i] = err1

            W[:, i1:i2] = Q1
            Losses += torch.sum(Losses1, 1) / 2

            W[:, i2:] -= Err1.matmul(Hinv[i1:i2, i2:])

        torch.cuda.synchronize()
        if isinstance(self.layer, transformers.Conv1D):
            W = W.t()
        self.layer.weight.data = W.reshape(self.layer.weight.shape).to(self.layer.weight.data.dtype)
//  [python]
def fasterprune(
        self, sparsity, prune_n=0, prune_m=0, blocksize=128, percdamp=.01
    ):
        W = self.layer.weight.data.clone()
        if isinstance(self.layer, nn.Conv2d):
            W = W.flatten(1)
        if isinstance(self.layer, transformers.Conv1D):
            W = W.t()
        W = W.float()

        tick = time.time()

        H = self.H
        del self.H
        dead = torch.diag(H) == 0
        H[dead, dead] = 1
        W[:, dead] = 0

        Losses = torch.zeros(self.rows, device=self.dev)

        damp = percdamp * torch.mean(torch.diag(H))
        diag = torch.arange(self.columns, device=self.dev)
        H[diag, diag] += damp
        H = torch.linalg.cholesky(H)
        H = torch.cholesky_inverse(H)
        H = torch.linalg.cholesky(H, upper=True)
        Hinv = H

        mask = None

        for i1 in range(0, self.columns, blocksize):
            i2 = min(i1 + blocksize, self.columns)
            count = i2 - i1

            W1 = W[:, i1:i2].clone()
            Q1 = torch.zeros_like(W1)
            Err1 = torch.zeros_like(W1)
            Losses1 = torch.zeros_like(W1)
            Hinv1 = Hinv[i1:i2, i1:i2]

            if prune_n == 0: 
                if mask is not None:
                    mask1 = mask[:, i1:i2]
                else:
                    tmp = W1 ** 2 / (torch.diag(Hinv1).reshape((1, -1))) ** 2
                    thresh = torch.sort(tmp.flatten())[0][int(tmp.numel() * sparsity)]
                    mask1 = tmp <= thresh
            else:
                mask1 = torch.zeros_like(W1) == 1

            for i in range(count):
                w = W1[:, i]
                d = Hinv1[i, i]

                if prune_n != 0 and i % prune_m == 0:
                    tmp = W1[:, i:(i + prune_m)] ** 2 / (torch.diag(Hinv1)[i:(i + prune_m)].reshape((1, -1))) ** 2
                    mask1.scatter_(1, i + torch.topk(tmp, prune_n, dim=1, largest=False)[1], True)

                q = w.clone()
                q[mask1[:, i]] = 0

                Q1[:, i] = q
                Losses1[:, i] = (w - q) ** 2 / d ** 2

                err1 = (w - q) / d 
                W1[:, i:] -= err1.unsqueeze(1).matmul(Hinv1[i, i:].unsqueeze(0))
                Err1[:, i] = err1

            W[:, i1:i2] = Q1
            Losses += torch.sum(Losses1, 1) / 2

            W[:, i2:] -= Err1.matmul(Hinv[i1:i2, i2:])

        torch.cuda.synchronize()
        if isinstance(self.layer, transformers.Conv1D):
            W = W.t()
        self.layer.weight.data = W.reshape(self.layer.weight.shape).to(self.layer.weight.data.dtype)
//  [python]
def hist_match(target: Tensor, source: Tensor, mode: str = "chol", eps: float = 1):
    target = target.permute(3, 0, 1, 2)
    source = source.permute(3, 0, 1, 2)
    c, b, h, w = target.shape

    if mode == "cdf":
        matched = cdf_match(target.view(c, -1), source.view(c, -1)).view(c, b, h, w)

    else:
        # based on https://github.com/ProGamerGov/Neural-Tools/blob/master/linear-color-transfer.py#L36

        mu_t = target.mean((2, 3), keepdim=True)
        hist_t = (target - mu_t).view(c, -1)  # [c, b * h * w]
        cov_t = hist_t @ hist_t.T / hist_t.shape[1] + eps * torch.eye(hist_t.shape[0], device=target.device)

        mu_s = source.mean((2, 3), keepdim=True)
        hist_s = (source - mu_s).view(c, -1)
        cov_s = hist_s @ hist_s.T / hist_s.shape[1] + eps * torch.eye(hist_s.shape[0], device=target.device)

        if mode == "chol":
            chol_t = torch.cholesky(cov_t)
            chol_s = torch.cholesky(cov_s)
            matched = chol_s @ torch.inverse(chol_t) @ hist_t

        elif mode == "pca":
            eva_t, eve_t = torch.linalg.eigh(cov_t, UPLO="U")
            Qt = eve_t @ torch.sqrt(torch.diag(eva_t)) @ eve_t.T
            eva_s, eve_s = torch.linalg.eigh(cov_s, UPLO="U")
            Qs = eve_s @ torch.sqrt(torch.diag(eva_s)) @ eve_s.T
            matched = Qs @ torch.inverse(Qt) @ hist_t

        else:  # mode == "sym"
            eva_t, eve_t = torch.linalg.eigh(cov_t, UPLO="U")
            Qt = eve_t @ torch.sqrt(torch.diag(eva_t)) @ eve_t.T
            Qt_Cs_Qt = Qt @ cov_s @ Qt
            eva_QtCsQt, eve_QtCsQt = torch.linalg.eigh(Qt_Cs_Qt, UPLO="U")
            QtCsQt = eve_QtCsQt @ torch.sqrt(torch.diag(eva_QtCsQt)) @ eve_QtCsQt.T
            matched = torch.inverse(Qt) @ QtCsQt @ torch.inverse(Qt) @ hist_t

        matched = matched.view(c, b, h, w) + mu_s

    return matched.permute(1, 2, 3, 0)
//  [python]
def hist_match(target: Tensor, source: Tensor, mode: str = "chol", eps: float = 1):
    target = target.permute(3, 0, 1, 2)
    source = source.permute(3, 0, 1, 2)
    c, b, h, w = target.shape

    if mode == "cdf":
        matched = cdf_match(target.view(c, -1), source.view(c, -1)).view(c, b, h, w)

    else:
        # based on https://github.com/ProGamerGov/Neural-Tools/blob/master/linear-color-transfer.py#L36

        mu_t = target.mean((2, 3), keepdim=True)
        hist_t = (target - mu_t).view(c, -1)  # [c, b * h * w]
        cov_t = hist_t @ hist_t.T / hist_t.shape[1] + eps * torch.eye(hist_t.shape[0], device=target.device)

        mu_s = source.mean((2, 3), keepdim=True)
        hist_s = (source - mu_s).view(c, -1)
        cov_s = hist_s @ hist_s.T / hist_s.shape[1] + eps * torch.eye(hist_s.shape[0], device=target.device)

        if mode == "chol":
            chol_t = torch.linalg.cholesky(cov_t)
            chol_s = torch.linalg.cholesky(cov_s)
            matched = chol_s @ torch.inverse(chol_t) @ hist_t

        elif mode == "pca":
            eva_t, eve_t = torch.linalg.eigh(cov_t, UPLO="U")
            Qt = eve_t @ torch.sqrt(torch.diag(eva_t)) @ eve_t.T
            eva_s, eve_s = torch.linalg.eigh(cov_s, UPLO="U")
            Qs = eve_s @ torch.sqrt(torch.diag(eva_s)) @ eve_s.T
            matched = Qs @ torch.inverse(Qt) @ hist_t

        else:  # mode == "sym"
            eva_t, eve_t = torch.linalg.eigh(cov_t, UPLO="U")
            Qt = eve_t @ torch.sqrt(torch.diag(eva_t)) @ eve_t.T
            Qt_Cs_Qt = Qt @ cov_s @ Qt
            eva_QtCsQt, eve_QtCsQt = torch.linalg.eigh(Qt_Cs_Qt, UPLO="U")
            QtCsQt = eve_QtCsQt @ torch.sqrt(torch.diag(eva_QtCsQt)) @ eve_QtCsQt.T
            matched = torch.inverse(Qt) @ QtCsQt @ torch.inverse(Qt) @ hist_t

        matched = matched.view(c, b, h, w) + mu_s

    return matched.permute(1, 2, 3, 0)
//  [python]
def test_torch_log_normal_spherical(benchmark: BenchmarkFixture):
    data = torch.randn(10000, 100)
    means = torch.randn(50, 100)
    covars = torch.rand(50)
    covar_matrices = torch.stack([torch.eye(means.size(-1)) * c for c in covars])

    cholesky = torch.cholesky(covar_matrices)
    distribution = MultivariateNormal(means, scale_tril=cholesky, validate_args=False)
    benchmark(distribution.log_prob, data.unsqueeze(1))
//  [python]
def test_torch_log_normal_spherical(benchmark: BenchmarkFixture):
    data = torch.randn(10000, 100)
    means = torch.randn(50, 100)
    covars = torch.rand(50)
    covar_matrices = torch.stack([torch.eye(means.size(-1)) * c for c in covars])

    cholesky = torch.linalg.cholesky(covar_matrices)
    distribution = MultivariateNormal(means, scale_tril=cholesky, validate_args=False)
    benchmark(distribution.log_prob, data.unsqueeze(1))
//  [python]
def test_torch_log_normal_diag(benchmark: BenchmarkFixture):
    data = torch.randn(10000, 100)
    means = torch.randn(50, 100)
    covars = torch.rand(50, 100)
    covar_matrices = torch.stack([torch.diag(c) for c in covars])

    cholesky = torch.cholesky(covar_matrices)
    distribution = MultivariateNormal(means, scale_tril=cholesky, validate_args=False)
    benchmark(distribution.log_prob, data.unsqueeze(1))
//  [python]
def test_torch_log_normal_diag(benchmark: BenchmarkFixture):
    data = torch.randn(10000, 100)
    means = torch.randn(50, 100)
    covars = torch.rand(50, 100)
    covar_matrices = torch.stack([torch.diag(c) for c in covars])

    cholesky = torch.linalg.cholesky(covar_matrices)
    distribution = MultivariateNormal(means, scale_tril=cholesky, validate_args=False)
    benchmark(distribution.log_prob, data.unsqueeze(1))
//  [python]
def test_torch_log_normal_full(benchmark: BenchmarkFixture):
    data = torch.randn(10000, 100)
    means = torch.randn(50, 100)
    A = torch.randn(50, 1000, 100)
    covars = A.permute(0, 2, 1).bmm(A)

    cholesky = torch.cholesky(covars)
    distribution = MultivariateNormal(means, scale_tril=cholesky, validate_args=False)
    benchmark(distribution.log_prob, data.unsqueeze(1))
//  [python]
def test_torch_log_normal_full(benchmark: BenchmarkFixture):
    data = torch.randn(10000, 100)
    means = torch.randn(50, 100)
    A = torch.randn(50, 1000, 100)
    covars = A.permute(0, 2, 1).bmm(A)

    cholesky = torch.linalg.cholesky(covars)
    distribution = MultivariateNormal(means, scale_tril=cholesky, validate_args=False)
    benchmark(distribution.log_prob, data.unsqueeze(1))
//  [python]
def fit(self, dataset: torch.utils.data.DataLoader) -> None:
        errors = []

        # Compute mean and covariance over the entire validation dataset
        counter = 0
        for b_inputs, b_targets in dataset:
            b_inputs = tuple(b_inp.to(self.dummy.device) for b_inp in b_inputs)
            b_targets = tuple(b_tar.to(self.dummy.device) for b_tar in b_targets)
            with torch.no_grad():
                pred = self.model(b_inputs)

            target, = b_targets

            error = target - pred
            for j in range(error.shape[0]):
                for j2 in range(error.shape[1]):
                    index = counter + j + j2
                    if len(errors) < index + 1:
                        errors.append([])
                    errors[counter + j + j2].append(error[j, j2])
            counter += error.shape[1]

        errors = errors[self.model.prediction_horizon - 1:-self.model.prediction_horizon + 1]

        errors = torch_utils.nested_list2tensor(errors)
        errors = errors.view(errors.shape[0], -1)
        mean = torch.mean(errors, dim=0)
        errors -= mean
        cov = torch.matmul(errors.T, errors)
        cov /= errors.shape[0] - 1

        # Add a small epsilon to the diagonal of the matrix to make it non-singular
        cov.diagonal().add_(1e-5)

        # This construction ensures that the resulting precision matrix is pos. semi-definite, even if the condition
        # number of the cov matrix is large
        cholesky = torch.cholesky(cov)
        precision = cov
        torch.cholesky_inverse(cholesky, out=precision)

        self.register_buffer('mean', mean)
        self.register_buffer('precision', precision)
//  [python]
def fit(self, dataset: torch.utils.data.DataLoader) -> None:
        errors = []

        # Compute mean and covariance over the entire validation dataset
        counter = 0
        for b_inputs, b_targets in dataset:
            b_inputs = tuple(b_inp.to(self.dummy.device) for b_inp in b_inputs)
            b_targets = tuple(b_tar.to(self.dummy.device) for b_tar in b_targets)
            with torch.no_grad():
                pred = self.model(b_inputs)

            target, = b_targets

            error = target - pred
            for j in range(error.shape[0]):
                for j2 in range(error.shape[1]):
                    index = counter + j + j2
                    if len(errors) < index + 1:
                        errors.append([])
                    errors[counter + j + j2].append(error[j, j2])
            counter += error.shape[1]

        errors = errors[self.model.prediction_horizon - 1:-self.model.prediction_horizon + 1]

        errors = torch_utils.nested_list2tensor(errors)
        errors = errors.view(errors.shape[0], -1)
        mean = torch.mean(errors, dim=0)
        errors -= mean
        cov = torch.matmul(errors.T, errors)
        cov /= errors.shape[0] - 1

        # Add a small epsilon to the diagonal of the matrix to make it non-singular
        cov.diagonal().add_(1e-5)

        # This construction ensures that the resulting precision matrix is pos. semi-definite, even if the condition
        # number of the cov matrix is large
        cholesky = torch.linalg.cholesky(cov)
        precision = cov
        torch.cholesky_inverse(cholesky, out=precision)

        self.register_buffer('mean', mean)
        self.register_buffer('precision', precision)
//  [python]
def test_torch_log_normal_tied(benchmark: BenchmarkFixture):
    data = torch.randn(10000, 100)
    means = torch.randn(50, 100)
    A = torch.randn(1000, 100)
    covars = A.t().mm(A)

    cholesky = torch.cholesky(covars)
    distribution = MultivariateNormal(means, scale_tril=cholesky, validate_args=False)
    benchmark(distribution.log_prob, data.unsqueeze(1))
//  [python]
def test_torch_log_normal_tied(benchmark: BenchmarkFixture):
    data = torch.randn(10000, 100)
    means = torch.randn(50, 100)
    A = torch.randn(1000, 100)
    covars = A.t().mm(A)

    cholesky = torch.linalg.cholesky(covars)
    distribution = MultivariateNormal(means, scale_tril=cholesky, validate_args=False)
    benchmark(distribution.log_prob, data.unsqueeze(1))
//  [python]
def forward(self, observations, add_bias=True):
        obs_dim = 2 * self.num_sensors
        bias_scale = self.log_bias_scale.exp()
        obs_noise = self.log_obs_noise.exp()
        trans_noise = self.log_trans_noise.exp()

        # bias distribution
        bias = Variable("bias", Reals[obs_dim])
        assert not torch.isnan(bias_scale), "bias scales was nan"
        bias_dist = dist_to_funsor(
            dist.MultivariateNormal(
                torch.zeros(obs_dim),
                scale_tril=bias_scale * torch.eye(2 * self.num_sensors),
            )
        )(value=bias)

        init_dist = dist.MultivariateNormal(
            torch.zeros(4), scale_tril=100.0 * torch.eye(4)
        )
        self.init = dist_to_funsor(init_dist)(value="state")

        # hidden states
        prev = Variable("prev", Reals[4])
        curr = Variable("curr", Reals[4])
        self.trans_dist = f_dist.MultivariateNormal(
            loc=prev @ NCV_TRANSITION_MATRIX,
            scale_tril=trans_noise * torch.cholesky(NCV_PROCESS_NOISE),
            value=curr,
        )

        state = Variable("state", Reals[4])
        obs = Variable("obs", Reals[obs_dim])
        observation_matrix = Tensor(
            torch.eye(4, 2)
            .unsqueeze(-1)
            .expand(-1, -1, self.num_sensors)
            .reshape(4, -1)
        )
        assert observation_matrix.output.shape == (
            4,
            obs_dim,
        ), observation_matrix.output.shape
        obs_loc = state @ observation_matrix
        if add_bias:
            obs_loc += bias
        self.observation_dist = f_dist.MultivariateNormal(
            loc=obs_loc, scale_tril=obs_noise * torch.eye(obs_dim), value=obs
        )

        logp = bias_dist
        curr = "state_init"
        logp += self.init(state=curr)
        for t, x in enumerate(observations):
            prev, curr = curr, "state_{}".format(t)
            logp += self.trans_dist(prev=prev, curr=curr)
            logp += self.observation_dist(state=curr, obs=x)
            # marginalize out previous state
            logp = logp.reduce(ops.logaddexp, prev)
        # marginalize out bias variable
        logp = logp.reduce(ops.logaddexp, "bias")

        # save posterior over the final state
        assert set(logp.inputs) == {"state_{}".format(len(observations) - 1)}
        posterior = funsor_to_mvn(logp, ndims=0)

        # marginalize out remaining variables
        logp = logp.reduce(ops.logaddexp)
        assert isinstance(logp, Tensor) and logp.shape == (), logp.pretty()
        return logp.data, posterior
//  [python]
def forward(self, observations, add_bias=True):
        obs_dim = 2 * self.num_sensors
        bias_scale = self.log_bias_scale.exp()
        obs_noise = self.log_obs_noise.exp()
        trans_noise = self.log_trans_noise.exp()

        # bias distribution
        bias = Variable("bias", Reals[obs_dim])
        assert not torch.isnan(bias_scale), "bias scales was nan"
        bias_dist = dist_to_funsor(
            dist.MultivariateNormal(
                torch.zeros(obs_dim),
                scale_tril=bias_scale * torch.eye(2 * self.num_sensors),
            )
        )(value=bias)

        init_dist = dist.MultivariateNormal(
            torch.zeros(4), scale_tril=100.0 * torch.eye(4)
        )
        self.init = dist_to_funsor(init_dist)(value="state")

        # hidden states
        prev = Variable("prev", Reals[4])
        curr = Variable("curr", Reals[4])
        self.trans_dist = f_dist.MultivariateNormal(
            loc=prev @ NCV_TRANSITION_MATRIX,
            scale_tril=trans_noise * torch.linalg.cholesky(NCV_PROCESS_NOISE),
            value=curr,
        )

        state = Variable("state", Reals[4])
        obs = Variable("obs", Reals[obs_dim])
        observation_matrix = Tensor(
            torch.eye(4, 2)
            .unsqueeze(-1)
            .expand(-1, -1, self.num_sensors)
            .reshape(4, -1)
        )
        assert observation_matrix.output.shape == (
            4,
            obs_dim,
        ), observation_matrix.output.shape
        obs_loc = state @ observation_matrix
        if add_bias:
            obs_loc += bias
        self.observation_dist = f_dist.MultivariateNormal(
            loc=obs_loc, scale_tril=obs_noise * torch.eye(obs_dim), value=obs
        )

        logp = bias_dist
        curr = "state_init"
        logp += self.init(state=curr)
        for t, x in enumerate(observations):
            prev, curr = curr, "state_{}".format(t)
            logp += self.trans_dist(prev=prev, curr=curr)
            logp += self.observation_dist(state=curr, obs=x)
            # marginalize out previous state
            logp = logp.reduce(ops.logaddexp, prev)
        # marginalize out bias variable
        logp = logp.reduce(ops.logaddexp, "bias")

        # save posterior over the final state
        assert set(logp.inputs) == {"state_{}".format(len(observations) - 1)}
        posterior = funsor_to_mvn(logp, ndims=0)

        # marginalize out remaining variables
        logp = logp.reduce(ops.logaddexp)
        assert isinstance(logp, Tensor) and logp.shape == (), logp.pretty()
        return logp.data, posterior
//  [python]
def cholesky_solve(A, b, eps=1e-8):
    diag = torch.diagonal(A)
    diag += eps
    if A.ndim > b.ndim:
        b = b.unsqueeze(dim=-1)
    u = torch.cholesky(A)
    return torch.cholesky_solve(b, u).squeeze(dim=-1)
//  [python]
def cholesky_solve(A, b, eps=1e-8):
    diag = torch.diagonal(A)
    diag += eps
    if A.ndim > b.ndim:
        b = b.unsqueeze(dim=-1)
    u = torch.linalg.cholesky(A)
    return torch.cholesky_solve(b, u).squeeze(dim=-1)
//  [python]
def fit(self, X, y):
        """should be called before forward() call.
        X: training input data point. N x D tensor for the data dimensionality D.
        y: training target data point. N x 1 tensor."""
        f = torch.zeros_like(y).float()
        N = X.shape[0]
        K = self.kernel_mat(X, X)
        while True:
            f = f.detach()
            W = -self.likelihood_func.get_hessian_of_log_likelihood(y, f)
            sqrt_W = W.sqrt()
            L = torch.cholesky(torch.eye(N, device=y.device) + sqrt_W.mm(K.mm(sqrt_W)))
            b = W.mm(f) + self.likelihood_func.get_jacobian_of_log_likelihood(y, f)
            a = b - sqrt_W.mm(torch.solve(L.T, torch.solve(L, sqrt_W.mm(K.mm(b)))[0]))
            diff = (torch.abs(K.mm(a) - f)).max()
            f = K.mm(a)
            if diff < self.eps:
                break

        approx_marginal_likelihood = (
            -0.5 * a.T.mm(f)
            - torch.log(torch.diag(L)).sum()
            + self.likelihood_func.get_log_likelihood(y, f)
        )
        self.X = X
        self.y = y
        self.sqrt_W = sqrt_W
        self.L = L
        self.K = K
        self.f = f
        return approx_marginal_likelihood
//  [python]
def fit(self, X, y):
        """should be called before forward() call.
        X: training input data point. N x D tensor for the data dimensionality D.
        y: training target data point. N x 1 tensor."""
        f = torch.zeros_like(y).float()
        N = X.shape[0]
        K = self.kernel_mat(X, X)
        while True:
            f = f.detach()
            W = -self.likelihood_func.get_hessian_of_log_likelihood(y, f)
            sqrt_W = W.sqrt()
            L = torch.linalg.cholesky(torch.eye(N, device=y.device) + sqrt_W.mm(K.mm(sqrt_W)))
            b = W.mm(f) + self.likelihood_func.get_jacobian_of_log_likelihood(y, f)
            a = b - sqrt_W.mm(torch.linalg.solve(L.T, torch.linalg.solve(L, sqrt_W.mm(K.mm(b)))))
            diff = (torch.abs(K.mm(a) - f)).max()
            f = K.mm(a)
            if diff < self.eps:
                break

        approx_marginal_likelihood = (
            -0.5 * a.T.mm(f)
            - torch.log(torch.diag(L)).sum()
            + self.likelihood_func.get_log_likelihood(y, f)
        )
        self.X = X
        self.y = y
        self.sqrt_W = sqrt_W
        self.L = L
        self.K = K
        self.f = f
        return approx_marginal_likelihood
//  [python]
def log_likelihood(self, x_train, y_train):
        K = self.kernel(x_train, x_train) + self.noise_variance.pow(2) * torch.eye(len(x_train))
        L = torch.cholesky(K)
        
        # y_dim = y_train.shape[1]
        # log_det_K = 2 * torch.sum(torch.log(torch.diag(L)))
        # Alpha = torch.cholesky_solve(y_train, L, upper = False)
        # return - 0.5 * ( (Alpha ** 2).sum() + log_det_K * y_dim + len(x_train) * y_dim * np.log(2 * np.pi) )
        
        # L = torch.linalg.cholesky(K)
        log_det_K = 2 * torch.sum(torch.log(torch.diag(L)))
        Alpha = torch.cholesky_solve(y_train, L, upper = False)
        
        # return - 0.5 * (Alpha.T @ Alpha + log_det_K + len(x_train) * np.log(2 * np.pi))
        return - 0.5 * ( (Alpha ** 2).sum() + log_det_K + len(x_train) * np.log(2 * np.pi))
//  [python]
def log_likelihood(self, x_train, y_train):
        K = self.kernel(x_train, x_train) + self.noise_variance.pow(2) * torch.eye(len(x_train))
        L = torch.linalg.cholesky(K)
        
        # y_dim = y_train.shape[1]
        # log_det_K = 2 * torch.sum(torch.log(torch.diag(L)))
        # Alpha = torch.cholesky_solve(y_train, L, upper = False)
        # return - 0.5 * ( (Alpha ** 2).sum() + log_det_K * y_dim + len(x_train) * y_dim * np.log(2 * np.pi) )
        
        # L = torch.linalg.cholesky(K)
        log_det_K = 2 * torch.sum(torch.log(torch.diag(L)))
        Alpha = torch.cholesky_solve(y_train, L, upper = False)
        
        # return - 0.5 * (Alpha.T @ Alpha + log_det_K + len(x_train) * np.log(2 * np.pi))
        return - 0.5 * ( (Alpha ** 2).sum() + log_det_K + len(x_train) * np.log(2 * np.pi))
//  [python]
def forward(self, x_train, y_train, x_test, Kinv_method='cholesky3'):
        """
        Computes the mean and covariance of the Gaussian process.

        Args:
            x_train (torch.Tensor): The input training data.
            y_train (torch.Tensor): The target training data.
            x_test (torch.Tensor): The input test data.
            Kinv_method (str, optional): The method used for computing the inverse of the covariance matrix. 
                Defaults to 'cholesky3'.

        Returns:
            tuple: A tuple containing the mean and covariance of the Gaussian process.

        Raises:
            ValueError: If Kinv_method is not 'direct' or 'cholesky'.

        """
        if isinstance(y_train, list):
            y_train_var = y_train[1]
            y_train = y_train[0]
        else:
            y_train_var = None
        K = self.kernel(x_train, x_train) + self.noise_variance.pow(2) * torch.eye(len(x_train))
        if y_train_var is not None:
            K = K + y_train_var
        K_s = self.kernel(x_train, x_test)
        K_ss = self.kernel(x_test, x_test)
        
        if Kinv_method == 'cholesky1':
            # kernel inverse is not stable, use cholesky decomposition instead
            L = torch.cholesky(K)
            L_inv = torch.inverse(L)
            K_inv = L_inv.T @ L_inv
            alpha = K_inv @ y_train
            mu = K_s.T @ alpha
            v = L_inv @ K_s
            var = K_ss - v.T @ v
        elif Kinv_method == 'cholesky3':
            # recommended implementation, fastest so far
            L = torch.cholesky(K)
            alpha = torch.cholesky_solve(y_train, L)
            mu = K_s.T @ alpha
            v = L.inverse() @ K_s
            var = K_ss - v.T @ v
        elif Kinv_method == 'direct':
            K_inv = torch.inverse(K)
            mu = K_s.T @ K_inv @ y_train
            var = K_ss - K_s.T @ K_inv @ K_s
        else:
            raise ValueError('Kinv_method should be either direct or cholesky')
        
        return mu.squeeze(), var
//  [python]
def forward(self, x_train, y_train, x_test, Kinv_method='cholesky3'):
        """
        Computes the mean and covariance of the Gaussian process.

        Args:
            x_train (torch.Tensor): The input training data.
            y_train (torch.Tensor): The target training data.
            x_test (torch.Tensor): The input test data.
            Kinv_method (str, optional): The method used for computing the inverse of the covariance matrix. 
                Defaults to 'cholesky3'.

        Returns:
            tuple: A tuple containing the mean and covariance of the Gaussian process.

        Raises:
            ValueError: If Kinv_method is not 'direct' or 'cholesky'.

        """
        if isinstance(y_train, list):
            y_train_var = y_train[1]
            y_train = y_train[0]
        else:
            y_train_var = None
        K = self.kernel(x_train, x_train) + self.noise_variance.pow(2) * torch.eye(len(x_train))
        if y_train_var is not None:
            K = K + y_train_var
        K_s = self.kernel(x_train, x_test)
        K_ss = self.kernel(x_test, x_test)
        
        if Kinv_method == 'cholesky1':
            # kernel inverse is not stable, use cholesky decomposition instead
            L = torch.linalg.cholesky(K)
            L_inv = torch.inverse(L)
            K_inv = L_inv.T @ L_inv
            alpha = K_inv @ y_train
            mu = K_s.T @ alpha
            v = L_inv @ K_s
            var = K_ss - v.T @ v
        elif Kinv_method == 'cholesky3':
            # recommended implementation, fastest so far
            L = torch.linalg.cholesky(K)
            alpha = torch.cholesky_solve(y_train, L)
            mu = K_s.T @ alpha
            v = L.inverse() @ K_s
            var = K_ss - v.T @ v
        elif Kinv_method == 'direct':
            K_inv = torch.inverse(K)
            mu = K_s.T @ K_inv @ y_train
            var = K_ss - K_s.T @ K_inv @ K_s
        else:
            raise ValueError('Kinv_method should be either direct or cholesky')
        
        return mu.squeeze(), var
//  [python]
def __cholesky(self):
        weight = self.weight.squeeze(dim=0)
        pos_semidef_matrix = torch.matmul(weight, weight.permute(0, 1, 3, 2))
        return torch.cholesky(pos_semidef_matrix + self.eps_eye)
//  [python]
def __cholesky(self):
        weight = self.weight.squeeze(dim=0)
        pos_semidef_matrix = torch.matmul(weight, weight.permute(0, 1, 3, 2))
        return torch.linalg.cholesky(pos_semidef_matrix + self.eps_eye)