jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def __init__(
        self,
        Qc_inv: Union[Variable, torch.Tensor],
        dt: Union[float, Variable, torch.Tensor],
        name: Optional[str] = None,
    ):
        super().__init__(name=name)
        dt = as_variable(dt)
        if dt.tensor.squeeze().ndim > 1:
            raise ValueError("dt must be a 0-D or 1-D tensor.")
        self.dt = dt
        self.dt.tensor = self.dt.tensor.view(-1, 1)
        if not (self.dt.tensor > 0).all():
            raise ValueError("dt must be greater than 0.")

        Qc_inv = as_variable(Qc_inv)
        if Qc_inv.ndim not in [2, 3]:
            raise ValueError("Qc_inv must be a single matrix or a batch of matrices.")
        if not Qc_inv.shape[-2] == Qc_inv.shape[-1]:
            raise ValueError("Qc_inv must contain square matrices.")
        self.Qc_inv = Qc_inv
        self.Qc_inv.tensor = (
            Qc_inv.tensor if Qc_inv.ndim == 3 else Qc_inv.tensor.unsqueeze(0)
        )
        try:
            torch.cholesky(Qc_inv.tensor)
        except RuntimeError:
            raise ValueError("Qc_inv must be positive definite.")

        self.register_aux_vars(["Qc_inv", "dt"])
//  [python]
def __init__(
        self,
        Qc_inv: Union[Variable, torch.Tensor],
        dt: Union[float, Variable, torch.Tensor],
        name: Optional[str] = None,
    ):
        super().__init__(name=name)
        dt = as_variable(dt)
        if dt.tensor.squeeze().ndim > 1:
            raise ValueError("dt must be a 0-D or 1-D tensor.")
        self.dt = dt
        self.dt.tensor = self.dt.tensor.view(-1, 1)
        if not (self.dt.tensor > 0).all():
            raise ValueError("dt must be greater than 0.")

        Qc_inv = as_variable(Qc_inv)
        if Qc_inv.ndim not in [2, 3]:
            raise ValueError("Qc_inv must be a single matrix or a batch of matrices.")
        if not Qc_inv.shape[-2] == Qc_inv.shape[-1]:
            raise ValueError("Qc_inv must contain square matrices.")
        self.Qc_inv = Qc_inv
        self.Qc_inv.tensor = (
            Qc_inv.tensor if Qc_inv.ndim == 3 else Qc_inv.tensor.unsqueeze(0)
        )
        try:
            torch.linalg.cholesky(Qc_inv.tensor)
        except RuntimeError:
            raise ValueError("Qc_inv must be positive definite.")

        self.register_aux_vars(["Qc_inv", "dt"])
//  [python]
def _compute_cost_weight(self) -> torch.Tensor:
        batch_size, dof, _ = self.Qc_inv.shape
        cost_weight = torch.empty(
            batch_size,
            2 * dof,
            2 * dof,
            dtype=self.Qc_inv.dtype,
            device=self.Qc_inv.device,
        )
        dt_data = self.dt.tensor.view(-1, 1, 1)
        Q11 = 12.0 * dt_data.pow(-3.0) * self.Qc_inv.tensor
        Q12 = -6.0 * dt_data.pow(-2.0) * self.Qc_inv.tensor
        Q22 = 4.0 * dt_data.reciprocal() * self.Qc_inv.tensor
        cost_weight[:, :dof, :dof] = Q11
        cost_weight[:, :dof:, dof:] = Q12
        cost_weight[:, dof:, :dof] = Q12
        cost_weight[:, dof:, dof:] = Q22
        return (
            torch.cholesky(cost_weight.transpose(-2, -1).conj())
            .transpose(-2, -1)
            .conj()
        )
//  [python]
def _compute_cost_weight(self) -> torch.Tensor:
        batch_size, dof, _ = self.Qc_inv.shape
        cost_weight = torch.empty(
            batch_size,
            2 * dof,
            2 * dof,
            dtype=self.Qc_inv.dtype,
            device=self.Qc_inv.device,
        )
        dt_data = self.dt.tensor.view(-1, 1, 1)
        Q11 = 12.0 * dt_data.pow(-3.0) * self.Qc_inv.tensor
        Q12 = -6.0 * dt_data.pow(-2.0) * self.Qc_inv.tensor
        Q22 = 4.0 * dt_data.reciprocal() * self.Qc_inv.tensor
        cost_weight[:, :dof, :dof] = Q11
        cost_weight[:, :dof:, dof:] = Q12
        cost_weight[:, dof:, :dof] = Q12
        cost_weight[:, dof:, dof:] = Q22
        return (
            torch.linalg.cholesky(cost_weight.transpose(-2, -1).conj())
            .transpose(-2, -1)
            .conj()
        )
//  [python]
def block_LDL(H, b, check_nan=True):
    n = H.shape[0]
    assert (n % b == 0)
    m = n // b
    try:
        L = torch.cholesky(H)
    except:
        return None
    DL = torch.diagonal(L.reshape(m, b, m, b), dim1=0, dim2=2).permute(2, 0, 1)
    D = (DL @ DL.permute(0, 2, 1)).cpu()
    DL = torch.linalg.inv(DL)
    L = L.view(n, m, b)
    for i in range(m):
        L[:, i, :] = L[:, i, :] @ DL[i, :, :]

    if check_nan and L.isnan().any():
        return None

    L = L.reshape(n, n)
    return (L, D.to(DL.device))
//  [python]
def block_LDL(H, b, check_nan=True):
    n = H.shape[0]
    assert (n % b == 0)
    m = n // b
    try:
        L = torch.linalg.cholesky(H)
    except:
        return None
    DL = torch.diagonal(L.reshape(m, b, m, b), dim1=0, dim2=2).permute(2, 0, 1)
    D = (DL @ DL.permute(0, 2, 1)).cpu()
    DL = torch.linalg.inv(DL)
    L = L.view(n, m, b)
    for i in range(m):
        L[:, i, :] = L[:, i, :] @ DL[i, :, :]

    if check_nan and L.isnan().any():
        return None

    L = L.reshape(n, n)
    return (L, D.to(DL.device))
//  [python]
def quantize(self):
        W = self.linear_module.weight.detach().to(dtype=torch.float, copy=True)

        scale, zero = self.find_params_weight(W)
        self.scales[:] = scale
        self.zeros[:] = zero

        H = self.H
        del self.H
        dead = torch.diag(H) == 0
        H[dead, dead] = 1
        W[:, dead] = 0
        if self.actorder:
            perm = torch.argsort(torch.diag(H), descending=True)
            W = W[:, perm]
            H = H[perm][:, perm]

        Losses = torch.zeros_like(W)
        Q = torch.zeros_like(W)

        damp = self.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

        for i1 in range(0, self.columns, self.blocksize):
            i2 = min(i1 + self.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]

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

                if self.groupsize != -1:
                    if (i1 + i) % self.groupsize == 0:
                        scale, zero = self.find_params_weight(
                            W[:, (i1 + i) : (i1 + i + self.groupsize)]
                        )
                        self.scales[:, (i1 + i) // self.groupsize] = scale
                        self.zeros[:, (i1 + i) // self.groupsize] = zero

                q = self.quantize_weight(w.unsqueeze(1), scale, zero, self.maxq)
                q = q.squeeze(1)
                assert q.dim() == 1
                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

            Q[:, i1:i2] = Q1
            Losses[:, i1:i2] = Losses1 / 2

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

        if self.actorder:
            invperm = torch.argsort(perm)
            Q = Q[:, invperm]

        weight = Q.reshape(self.linear_module.weight.shape).to(
            self.linear_module.weight.data.dtype
        )
        error = torch.sum(Losses).item()

        q_module = ColBlockQuantizedLinear(
            self.linear_module.in_features,
            self.linear_module.out_features,
            self.linear_module.bias is not None,
            bits=self.bits,
            tile_cols=self.groupsize,
        ).to(self.dev)
        q_module.scales = self.scales
        q_module.zeros = self.zeros
        q_module.pack_weight(weight)
        q_module.bias = self.linear_module.bias
        return q_module, error
//  [python]
def quantize(self):
        W = self.linear_module.weight.detach().to(dtype=torch.float, copy=True)

        scale, zero = self.find_params_weight(W)
        self.scales[:] = scale
        self.zeros[:] = zero

        H = self.H
        del self.H
        dead = torch.diag(H) == 0
        H[dead, dead] = 1
        W[:, dead] = 0
        if self.actorder:
            perm = torch.argsort(torch.diag(H), descending=True)
            W = W[:, perm]
            H = H[perm][:, perm]

        Losses = torch.zeros_like(W)
        Q = torch.zeros_like(W)

        damp = self.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

        for i1 in range(0, self.columns, self.blocksize):
            i2 = min(i1 + self.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]

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

                if self.groupsize != -1:
                    if (i1 + i) % self.groupsize == 0:
                        scale, zero = self.find_params_weight(
                            W[:, (i1 + i) : (i1 + i + self.groupsize)]
                        )
                        self.scales[:, (i1 + i) // self.groupsize] = scale
                        self.zeros[:, (i1 + i) // self.groupsize] = zero

                q = self.quantize_weight(w.unsqueeze(1), scale, zero, self.maxq)
                q = q.squeeze(1)
                assert q.dim() == 1
                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

            Q[:, i1:i2] = Q1
            Losses[:, i1:i2] = Losses1 / 2

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

        if self.actorder:
            invperm = torch.argsort(perm)
            Q = Q[:, invperm]

        weight = Q.reshape(self.linear_module.weight.shape).to(
            self.linear_module.weight.data.dtype
        )
        error = torch.sum(Losses).item()

        q_module = ColBlockQuantizedLinear(
            self.linear_module.in_features,
            self.linear_module.out_features,
            self.linear_module.bias is not None,
            bits=self.bits,
            tile_cols=self.groupsize,
        ).to(self.dev)
        q_module.scales = self.scales
        q_module.zeros = self.zeros
        q_module.pack_weight(weight)
        q_module.bias = self.linear_module.bias
        return q_module, error
//  [python]
def forward(self, x, L, mass, evals, evecs):

        # project times to the positive halfspace
        # (and away from 0 in the incredibly rare chance that they get stuck)
        with torch.no_grad():
            self.diffusion_time.data = torch.clamp(self.diffusion_time, min=1e-8)

        if x.shape[-1] != self.C_inout:
            raise ValueError("Tensor has wrong shape = {}. Last dim shape should have number of channels = {}".format(
                x.shape, self.C_inout))

        if self.method == 'spectral':

            # Transform to spectral
            x_spec = to_basis(x, evecs, mass)

            # Diffuse
            time = self.diffusion_time
            diffusion_coefs = torch.exp(-evals.unsqueeze(-1) * time.unsqueeze(0))
            x_diffuse_spec = diffusion_coefs * x_spec

            # Transform back to per-vertex
            x_diffuse = from_basis(x_diffuse_spec, evecs)

        elif self.method == 'implicit_dense':
            V = x.shape[-2]

            # Form the dense matrices (M + tL) with dims (B,C,V,V)
            mat_dense = L.to_dense().unsqueeze(1).expand(-1, self.C_inout, V, V).clone()
            mat_dense *= self.diffusion_time.unsqueeze(0).unsqueeze(-1).unsqueeze(-1)
            mat_dense += torch.diag_embed(mass).unsqueeze(1)

            # Factor the system
            cholesky_factors = torch.cholesky(mat_dense)

            # Solve the system
            rhs = x * mass.unsqueeze(-1)
            rhsT = torch.transpose(rhs, 1, 2).unsqueeze(-1)
            sols = torch.cholesky_solve(rhsT, cholesky_factors)
            x_diffuse = torch.transpose(sols.squeeze(-1), 1, 2)

        else:
            raise ValueError("unrecognized method")

        return x_diffuse
//  [python]
def forward(self, x, L, mass, evals, evecs):

        # project times to the positive halfspace
        # (and away from 0 in the incredibly rare chance that they get stuck)
        with torch.no_grad():
            self.diffusion_time.data = torch.clamp(self.diffusion_time, min=1e-8)

        if x.shape[-1] != self.C_inout:
            raise ValueError("Tensor has wrong shape = {}. Last dim shape should have number of channels = {}".format(
                x.shape, self.C_inout))

        if self.method == 'spectral':

            # Transform to spectral
            x_spec = to_basis(x, evecs, mass)

            # Diffuse
            time = self.diffusion_time
            diffusion_coefs = torch.exp(-evals.unsqueeze(-1) * time.unsqueeze(0))
            x_diffuse_spec = diffusion_coefs * x_spec

            # Transform back to per-vertex
            x_diffuse = from_basis(x_diffuse_spec, evecs)

        elif self.method == 'implicit_dense':
            V = x.shape[-2]

            # Form the dense matrices (M + tL) with dims (B,C,V,V)
            mat_dense = L.to_dense().unsqueeze(1).expand(-1, self.C_inout, V, V).clone()
            mat_dense *= self.diffusion_time.unsqueeze(0).unsqueeze(-1).unsqueeze(-1)
            mat_dense += torch.diag_embed(mass).unsqueeze(1)

            # Factor the system
            cholesky_factors = torch.linalg.cholesky(mat_dense)

            # Solve the system
            rhs = x * mass.unsqueeze(-1)
            rhsT = torch.transpose(rhs, 1, 2).unsqueeze(-1)
            sols = torch.cholesky_solve(rhsT, cholesky_factors)
            x_diffuse = torch.transpose(sols.squeeze(-1), 1, 2)

        else:
            raise ValueError("unrecognized method")

        return x_diffuse
//  [python]
def logdet(self):
        if self.param_type == "scalar":
            return 2 * self.x_dim * torch.log(self.post_process(self.q, self.param_type))
        elif self.param_type == "diag":
            return 2 * self.post_process(self.q, self.param_type).log().sum()
        elif self.param_type == "tril":
            return 2 * self.post_process(self.q, self.param_type).diagonal(dim1=-2,dim2=-1).log().sum(-1)
        elif self.param_type == "full":
            return 2 * torch.cholesky(self.q).diagonal(dim1=-2,dim2=-1).log().sum(-1)
//  [python]
def logdet(self):
        if self.param_type == "scalar":
            return 2 * self.x_dim * torch.log(self.post_process(self.q, self.param_type))
        elif self.param_type == "diag":
            return 2 * self.post_process(self.q, self.param_type).log().sum()
        elif self.param_type == "tril":
            return 2 * self.post_process(self.q, self.param_type).diagonal(dim1=-2,dim2=-1).log().sum(-1)
        elif self.param_type == "full":
            return 2 * torch.linalg.cholesky(self.q).diagonal(dim1=-2,dim2=-1).log().sum(-1)
//  [python]
def chol(self):
        if self.param_type == "scalar":
            if self.q_shape is None:
                return self.post_process(self.q, self.param_type) * torch.eye(self.x_dim, device=self.q.device)
            else:
                return self.post_process(self.q, self.param_type) * torch.cholesky(
                    torch.exp(self.q_shape * self.base))
        elif self.param_type == "diag":
            return self.post_process(self.q, self.param_type) * torch.eye(self.x_dim, device=self.q.device)
        elif self.param_type == "tril":
            return self.post_process(self.q, self.param_type)
        elif self.param_type == "full":
            return torch.cholesky(self.q)
//  [python]
def chol(self):
        if self.param_type == "scalar":
            if self.q_shape is None:
                return self.post_process(self.q, self.param_type) * torch.eye(self.x_dim, device=self.q.device)
            else:
                return self.post_process(self.q, self.param_type) * torch.linalg.cholesky(
                    torch.exp(self.q_shape * self.base))
        elif self.param_type == "diag":
            return self.post_process(self.q, self.param_type) * torch.eye(self.x_dim, device=self.q.device)
        elif self.param_type == "tril":
            return self.post_process(self.q, self.param_type)
        elif self.param_type == "full":
            return torch.linalg.cholesky(self.q)
//  [python]
def forward(self, X):
        if self.param_type == "scalar":
            if self.q_shape is None:
                X = X + self.post_process(self.q, self.param_type) * torch.randn_like(X)
            else:
                chol = self.post_process(self.q, self.param_type) * torch.cholesky(
                    torch.exp(self.q_shape * self.base))
                X = X + torch.randn_like(X) @ chol.t()
        elif self.param_type == "diag":
            X = X + self.post_process(self.q, self.param_type) * torch.randn_like(X)  # (x_dim) * (*bs, N_ensem, x_dim)
        elif self.param_type == "tril":
            # batch_shape = X.shape[:-1]
            chol = self.post_process(self.q, self.param_type)
            X = X + torch.randn_like(X) @ chol.t()  # (*bs, N_ensem, x_dim) @ (x_dim, x_dim)
            # X = X + torch.distributions.MultivariateNormal(torch.zeros(self.x_dim, device=self.q.device), scale_tril=chol).sample(batch_shape)
        elif self.param_type == "full":
            batch_shape = X.shape[:-1]
            chol = torch.cholesky(self.q)
            X = X + torch.distributions.MultivariateNormal(torch.zeros(self.x_dim, device=self.q.device),
                                                           scale_tril=chol).sample(batch_shape)  # (*bs, N_ensem, x_dim)
        return X
//  [python]
def forward(self, X):
        if self.param_type == "scalar":
            if self.q_shape is None:
                X = X + self.post_process(self.q, self.param_type) * torch.randn_like(X)
            else:
                chol = self.post_process(self.q, self.param_type) * torch.linalg.cholesky(
                    torch.exp(self.q_shape * self.base))
                X = X + torch.randn_like(X) @ chol.t()
        elif self.param_type == "diag":
            X = X + self.post_process(self.q, self.param_type) * torch.randn_like(X)  # (x_dim) * (*bs, N_ensem, x_dim)
        elif self.param_type == "tril":
            # batch_shape = X.shape[:-1]
            chol = self.post_process(self.q, self.param_type)
            X = X + torch.randn_like(X) @ chol.t()  # (*bs, N_ensem, x_dim) @ (x_dim, x_dim)
            # X = X + torch.distributions.MultivariateNormal(torch.zeros(self.x_dim, device=self.q.device), scale_tril=chol).sample(batch_shape)
        elif self.param_type == "full":
            batch_shape = X.shape[:-1]
            chol = torch.linalg.cholesky(self.q)
            X = X + torch.distributions.MultivariateNormal(torch.zeros(self.x_dim, device=self.q.device),
                                                           scale_tril=chol).sample(batch_shape)  # (*bs, N_ensem, x_dim)
        return X
//  [python]
def _cholesky_solve(A, b, eps=1e-8):
    A = _add_value_to_diagonal(A, 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 forward(ctx, H, b):
        # don't crash training if cholesky decomp fails
        try:
            U = torch.linalg.cholesky(H)
            xs = torch.cholesky_solve(b, U)
            ctx.save_for_backward(U, xs)
            ctx.failed = False
        except Exception as e:
            print(e)
            ctx.failed = True
            xs = torch.zeros_like(b)

        return xs
//  [python]
def robust_cholesky(matrix: torch.Tensor) -> torch.Tensor:
    """
    Implements a Cholesky decomposition.
    If the Cholesky decomposition fails, it is retried with increasing added jitter on the diagonal.
    :param matrix: Symmetric positive semi-definite matrix to factorize.
    :return: Approximate cholesky factor L such that (approximately) LL^T = matrix
    """
    eps = 1e-5 * matrix.trace() / matrix.shape[-1] + 1e-30
    for i in range(15):
        try:
            # on the CPU, cholesky can fail by throwing RuntimeError
            # on the GPU, cholesky can fail by returning NaN values
            # we catch both cases
            L = torch.cholesky(matrix)
            if not L.isnan().any().item():
                return L
        except RuntimeError:
            pass

        print('Increasing jitter for Cholesky decomposition', flush=True)
        matrix += eps * torch.eye(matrix.shape[-1], device=matrix.device, dtype=matrix.dtype)
        eps *= 2
    raise RuntimeError('Could not Cholesky decompose the matrix')
//  [python]
def robust_cholesky(matrix: torch.Tensor) -> torch.Tensor:
    """
    Implements a Cholesky decomposition.
    If the Cholesky decomposition fails, it is retried with increasing added jitter on the diagonal.
    :param matrix: Symmetric positive semi-definite matrix to factorize.
    :return: Approximate cholesky factor L such that (approximately) LL^T = matrix
    """
    eps = 1e-5 * matrix.trace() / matrix.shape[-1] + 1e-30
    for i in range(15):
        try:
            # on the CPU, cholesky can fail by throwing RuntimeError
            # on the GPU, cholesky can fail by returning NaN values
            # we catch both cases
            L = torch.linalg.cholesky(matrix)
            if not L.isnan().any().item():
                return L
        except RuntimeError:
            pass

        print('Increasing jitter for Cholesky decomposition', flush=True)
        matrix += eps * torch.eye(matrix.shape[-1], device=matrix.device, dtype=matrix.dtype)
        eps *= 2
    raise RuntimeError('Could not Cholesky decompose the matrix')
//  [python]
def gmm(cfg,all_fea, pi, mu, all_output):    
    Cov = []
    dist = []
    log_probs = []
    
    for i in range(len(mu)):
        temp = all_fea - mu[i]
        predi = all_output[:,i].unsqueeze(dim=-1)
        Covi = torch.matmul(temp.t(), temp * predi.expand_as(temp)) / (predi.sum()) + cfg.epsilon * torch.eye(temp.shape[1]).cuda()
        try:
            chol = torch.cholesky(Covi)
        except RuntimeError:
            Covi += cfg.epsilon * torch.eye(temp.shape[1]).cuda() * 100
            chol = torch.cholesky(Covi)
        chol_inv = torch.inverse(chol)
        Covi_inv = torch.matmul(chol_inv.t(), chol_inv)
        logdet = torch.logdet(Covi)
        mah_dist = (torch.matmul(temp, Covi_inv) * temp).sum(dim=1)
        log_prob = -0.5*(Covi.shape[0] * np.log(2*math.pi) + logdet + mah_dist) + torch.log(pi)[i]
        Cov.append(Covi)
        log_probs.append(log_prob)
        dist.append(mah_dist)
    Cov = torch.stack(Cov, dim=0)
    dist = torch.stack(dist, dim=0).t()
    log_probs = torch.stack(log_probs, dim=0).t()
    zz = log_probs - torch.logsumexp(log_probs, dim=1, keepdim=True).expand_as(log_probs)
    gamma = torch.exp(zz)
    
    return zz, gamma
//  [python]
def gmm(cfg,all_fea, pi, mu, all_output):    
    Cov = []
    dist = []
    log_probs = []
    
    for i in range(len(mu)):
        temp = all_fea - mu[i]
        predi = all_output[:,i].unsqueeze(dim=-1)
        Covi = torch.matmul(temp.t(), temp * predi.expand_as(temp)) / (predi.sum()) + cfg.epsilon * torch.eye(temp.shape[1]).cuda()
        try:
            chol = torch.linalg.cholesky(Covi)
        except RuntimeError:
            Covi += cfg.epsilon * torch.eye(temp.shape[1]).cuda() * 100
            chol = torch.linalg.cholesky(Covi)
        chol_inv = torch.inverse(chol)
        Covi_inv = torch.matmul(chol_inv.t(), chol_inv)
        logdet = torch.logdet(Covi)
        mah_dist = (torch.matmul(temp, Covi_inv) * temp).sum(dim=1)
        log_prob = -0.5*(Covi.shape[0] * np.log(2*math.pi) + logdet + mah_dist) + torch.log(pi)[i]
        Cov.append(Covi)
        log_probs.append(log_prob)
        dist.append(mah_dist)
    Cov = torch.stack(Cov, dim=0)
    dist = torch.stack(dist, dim=0).t()
    log_probs = torch.stack(log_probs, dim=0).t()
    zz = log_probs - torch.logsumexp(log_probs, dim=1, keepdim=True).expand_as(log_probs)
    gamma = torch.exp(zz)
    
    return zz, gamma
//  [python]
def cholesky_update(self, constmean: float, log_amp: float, log_beta: torch.FloatTensor,
                        log_likelihood_noise: float):
        """
        Function used to update the Cholesky matrix
        :param constmean
        :param log_amp:
        :param log_beta:
        :param log_likelihood_noise:
        :return:
        """
        self.gram_mat_update(constmean, log_amp, log_beta, log_likelihood_noise)

        eye_mat = torch.diag(self.gram_mat.new_ones(self.gram_mat.size(0)))
        for jitter_const in [0, 1e-8, 1e-7, 1e-6, 1e-5, 1e-4, 1e-3, 1e-2]:
            chol_jitter = torch.trace(self.gram_mat).item() * jitter_const
            try:
                # cholesky is lower triangular matrix
                self.cholesky = torch.cholesky(self.gram_mat + eye_mat * chol_jitter)
                self.jitter = chol_jitter
                return
            except RuntimeError:
                pass
        raise RuntimeError('Absolute entry values of Gram matrix are between %.4E~%.4E with trace %.4E' %
                           (torch.min(torch.abs(self.gram_mat)).item(), torch.max(torch.abs(self.gram_mat)).item(),
                            torch.trace(self.gram_mat).item()))
//  [python]
def cholesky_update(self, constmean: float, log_amp: float, log_beta: torch.FloatTensor,
                        log_likelihood_noise: float):
        """
        Function used to update the Cholesky matrix
        :param constmean
        :param log_amp:
        :param log_beta:
        :param log_likelihood_noise:
        :return:
        """
        self.gram_mat_update(constmean, log_amp, log_beta, log_likelihood_noise)

        eye_mat = torch.diag(self.gram_mat.new_ones(self.gram_mat.size(0)))
        for jitter_const in [0, 1e-8, 1e-7, 1e-6, 1e-5, 1e-4, 1e-3, 1e-2]:
            chol_jitter = torch.trace(self.gram_mat).item() * jitter_const
            try:
                # cholesky is lower triangular matrix
                self.cholesky = torch.linalg.cholesky(self.gram_mat + eye_mat * chol_jitter)
                self.jitter = chol_jitter
                return
            except RuntimeError:
                pass
        raise RuntimeError('Absolute entry values of Gram matrix are between %.4E~%.4E with trace %.4E' %
                           (torch.min(torch.abs(self.gram_mat)).item(), torch.max(torch.abs(self.gram_mat)).item(),
                            torch.trace(self.gram_mat).item()))
//  [python]
def forward(self, mu, sigma, pi, x, mask, is_tril=False):
        """
        N = batch size
        T = number of time steps
        k = number of states
        m = number of mixtures

        mu.shape (N, T, m, k)
        sigma.shape (N, T, m, k, k)
        pi.shape (N, m)
        x.shape (N, T, k)
        mask.shape (N, T)
        """

        pi_expanded = pi.unsqueeze(1).expand(-1, mu.shape[1], -1)
        mix = tdist.Categorical(logits=pi_expanded)  # Broadcast this over time dim.
        L = sigma if is_tril else torch.cholesky(sigma)
        mvn = tdist.MultivariateNormal(mu, scale_tril=L)
        gmm = tdist.MixtureSameFamily(mix, mvn)
        comp_loss = gmm.log_prob(x).neg()  # (N, T)
        return mask_reduce_losses(comp_loss, mask)
//  [python]
def forward(self, mu, sigma, pi, x, mask, is_tril=False):
        """
        N = batch size
        T = number of time steps
        k = number of states
        m = number of mixtures

        mu.shape (N, T, m, k)
        sigma.shape (N, T, m, k, k)
        pi.shape (N, m)
        x.shape (N, T, k)
        mask.shape (N, T)
        """

        pi_expanded = pi.unsqueeze(1).expand(-1, mu.shape[1], -1)
        mix = tdist.Categorical(logits=pi_expanded)  # Broadcast this over time dim.
        L = sigma if is_tril else torch.linalg.cholesky(sigma)
        mvn = tdist.MultivariateNormal(mu, scale_tril=L)
        gmm = tdist.MixtureSameFamily(mix, mvn)
        comp_loss = gmm.log_prob(x).neg()  # (N, T)
        return mask_reduce_losses(comp_loss, mask)
//  [python]
def predict(
        x: Tensor,
        V: Tensor,
        A: Tensor,
        Q: Tensor,
        H: Tensor,
        R: Tensor,
        cholesky=True,
) -> Tuple:
    """
    x(t) | x(t-1), P(t-1), u(t), A, B, Q
    :param x: previous state, (xdim, batch)
    :param V: previous posteriori covariance, (xdim, xdim)
    :param A: transition matrix, (xdim, xdim)
    :param Q: state noise covariance, (xdim, xdim)
    :param H: observation matrix, (ydim, xdim)
    :param R: observation noise covariance, (ydim, ydim)
    :param cholesky: True if V is Cholesky form, default=True
    :return:
        yhat: predicted observation, (ydim, batch)
        xhat: predicted mean, (xdim, batch)
        Vhat: predicted covariance or its Choleksy, (xdim, xdim)
    """
    n_sample = H.shape[0]
    xhat = A.mm(x)  # Ax
    if cholesky:
        L = V
    else:
        L = t.cholesky(V)
    AL = A.mm(L)
    Vhat = AL.mm(AL.t()) + Q  # APA' + Q, n samples one step equivalent to one sample n steps
    yhat = H.mm(xhat)
    if cholesky:
        Vhat = t.cholesky(Vhat)
    return yhat, xhat, Vhat
//  [python]
def predict(
        x: Tensor,
        V: Tensor,
        A: Tensor,
        Q: Tensor,
        H: Tensor,
        R: Tensor,
        cholesky=True,
) -> Tuple:
    """
    x(t) | x(t-1), P(t-1), u(t), A, B, Q
    :param x: previous state, (xdim, batch)
    :param V: previous posteriori covariance, (xdim, xdim)
    :param A: transition matrix, (xdim, xdim)
    :param Q: state noise covariance, (xdim, xdim)
    :param H: observation matrix, (ydim, xdim)
    :param R: observation noise covariance, (ydim, ydim)
    :param cholesky: True if V is Cholesky form, default=True
    :return:
        yhat: predicted observation, (ydim, batch)
        xhat: predicted mean, (xdim, batch)
        Vhat: predicted covariance or its Choleksy, (xdim, xdim)
    """
    n_sample = H.shape[0]
    xhat = A.mm(x)  # Ax
    if cholesky:
        L = V
    else:
        L = linalg.cholesky(V)
    AL = A.mm(L)
    Vhat = AL.mm(AL.t()) + Q  # APA' + Q, n samples one step equivalent to one sample n steps
    yhat = H.mm(xhat)
    if cholesky:
        Vhat = linalg.cholesky(Vhat)
    return yhat, xhat, Vhat
//  [python]
def model(self):
        self.set_mode("model")

        N = self.X.size(0)
        Kff = self.kernel(self.X)
        Kff = Kff + self.noise
        Kff.view(-1)[:: N + 1] += self.jitter 
#         Kff.view(-1)[:: N + 1] += self.jitter + self.noise  # add noise to diagonal

        Lff = torch.cholesky(Kff)

        zero_loc = self.X.new_zeros(self.X.size(0))
        f_loc = zero_loc + self.mean_function(self.X)
        if self.y is None:
            f_var = Lff.pow(2).sum(dim=-1)
            return f_loc, f_var
        else:
            return pyro.sample(
                self._pyro_get_fullname("y"),
                dist.MultivariateNormal(f_loc, scale_tril=Lff)
                .expand_by(self.y.shape[:-1])
                .to_event(self.y.dim() - 1),
                obs=self.y,
            )
//  [python]
def model(self):
        self.set_mode("model")

        N = self.X.size(0)
        Kff = self.kernel(self.X)
        Kff = Kff + self.noise
        Kff.view(-1)[:: N + 1] += self.jitter 
#         Kff.view(-1)[:: N + 1] += self.jitter + self.noise  # add noise to diagonal

        Lff = torch.linalg.cholesky(Kff)

        zero_loc = self.X.new_zeros(self.X.size(0))
        f_loc = zero_loc + self.mean_function(self.X)
        if self.y is None:
            f_var = Lff.pow(2).sum(dim=-1)
            return f_loc, f_var
        else:
            return pyro.sample(
                self._pyro_get_fullname("y"),
                dist.MultivariateNormal(f_loc, scale_tril=Lff)
                .expand_by(self.y.shape[:-1])
                .to_event(self.y.dim() - 1),
                obs=self.y,
            )
//  [python]
def update(y: Tensor,
           yhat: Tensor,
           xhat: Tensor,
           Vhat: Tensor,
           H: Tensor,
           R: Tensor,
           cholesky=True,
           ) -> Tuple:
    """
    :param y: measurement, (ydim,)
    :param yhat: predicted measurement, (ydim,)
    :param xhat: predicted state, (xdim,)
    :param Vhat: predicted covariance, (xdim, xdim)
    :param H: measurement matrix, (ydim, xdim)
    :param R: measurement noise covariance, (ydim, ydim)
    :param cholesky: True if Vhat is Cholesky form, default=True
    :return:
        x: posterior mean, (xdim, batch)
        V: posterior covariance or its Cholesky, (xdim, xdim)
    """
    e = y - yhat
    if cholesky:
        Lhat = Vhat
        Vhat = Lhat.mm(Lhat.t())
    else:
        Lhat = t.cholesky(Vhat)
    HL = H.mm(Lhat)
    S = HL.mm(HL.t()) + R  # HVH' + R

    L = t.cholesky(S)
    # K = H.cholesky_solve(L).mm(Vhat)  # L^{-1}HV
    # K = H.mm(Vhat).cholesky_solve(L)  # L^{-1}HV
    G = H.mm(Vhat).triangular_solve(L, upper=False).solution.t()
    # G' = L^{-1}HV, gain K = VH'S^{-1} = VH'(LL')^{-1} = VH'L'^{-1}L^{-1} = G L^{-1}

    x = xhat + G.mm(e.triangular_solve(L, upper=False).solution)
    V = Vhat - G.mm(G.t())  # minus is dangerous
    # V = positivize(V)
    # assert symmetric(V), 'V is asymmetric'
    if cholesky:
        try:
            V = t.cholesky(V)
        except RuntimeError:
            print('Singular covariance', t.linalg.eigvalsh(V))

    return x, V
//  [python]
def update(y: Tensor,
           yhat: Tensor,
           xhat: Tensor,
           Vhat: Tensor,
           H: Tensor,
           R: Tensor,
           cholesky=True,
           ) -> Tuple:
    """
    :param y: measurement, (ydim,)
    :param yhat: predicted measurement, (ydim,)
    :param xhat: predicted state, (xdim,)
    :param Vhat: predicted covariance, (xdim, xdim)
    :param H: measurement matrix, (ydim, xdim)
    :param R: measurement noise covariance, (ydim, ydim)
    :param cholesky: True if Vhat is Cholesky form, default=True
    :return:
        x: posterior mean, (xdim, batch)
        V: posterior covariance or its Cholesky, (xdim, xdim)
    """
    e = y - yhat
    if cholesky:
        Lhat = Vhat
        Vhat = Lhat.mm(Lhat.t())
    else:
        Lhat = linalg.cholesky(Vhat)
    HL = H.mm(Lhat)
    S = HL.mm(HL.t()) + R  # HVH' + R

    L = linalg.cholesky(S)
    # K = H.cholesky_solve(L).mm(Vhat)  # L^{-1}HV
    # K = H.mm(Vhat).cholesky_solve(L)  # L^{-1}HV
    G = H.mm(Vhat).triangular_solve(L, upper=False).solution.t()
    # G' = L^{-1}HV, gain K = VH'S^{-1} = VH'(LL')^{-1} = VH'L'^{-1}L^{-1} = G L^{-1}

    x = xhat + G.mm(e.triangular_solve(L, upper=False).solution)
    V = Vhat - G.mm(G.t())  # minus is dangerous
    # V = positivize(V)
    # assert symmetric(V), 'V is asymmetric'
    if cholesky:
        try:
            V = linalg.cholesky(V)
        except RuntimeError:
            print('Singular covariance', torch.linalg.eigvalsh(V))

    return x, V
//  [python]
def joseph_update(y: Tensor,
                  yhat: Tensor,
                  xhat: Tensor,
                  Vhat: Tensor,
                  H: Tensor,
                  R: Tensor,
                  cholesky=True,
                  ) -> Tuple:
    """
    :param y: measurement, (ydim,)
    :param yhat: predicted measurement, (ydim,)
    :param xhat: predicted state, (xdim,)
    :param Vhat: predicted covariance, (xdim, xdim)
    :param H: measurement matrix, (ydim, xdim)
    :param R: measurement noise covariance, (ydim, ydim)
    :param cholesky: True if Vhat is Cholesky form, default=True
    :return:
        x: posterior mean, (xdim, batch)
        V: posterior covariance or its Cholesky, (xdim, xdim)
    """
    e = y - yhat
    if cholesky:
        Lhat = Vhat
        Vhat = Lhat.mm(Lhat.t())
    else:
        Lhat = torch.cholesky(Vhat)
    HL = H.mm(Lhat)
    S = HL.mm(HL.t()) + R  # HVH' + R

    L = torch.cholesky(S)
    G = H.mm(Vhat).cholesky_solve(L).t()  # L^{-1}HV, gain K = VH'S^{-1} = VH'(LL')^{-1} = VH'L'^{-1}L^{-1} = G L^{-1}
    x = xhat + G.mm(e.cholesky_solve(L))
    # V = (I - KH) Vhat (I - KH)' + K R K'
    eye = torch.eye(Vhat.shape[0])
    IminusKH = eye - G.mm(H.cholesky_solve(L))
    IminusKHLhat = IminusKH.mm(Lhat)
    KR = G.mm(R.sqrt().cholesky_solve(L))  # R's supposed to be diagonal
    V = IminusKHLhat.mm(IminusKHLhat.t()) + KR.mm(KR.t())
    if cholesky:
        V = torch.cholesky(V)
        # print(torch.linalg.eigvalsh(V))

    return x, V
//  [python]
def joseph_update(y: Tensor,
                  yhat: Tensor,
                  xhat: Tensor,
                  Vhat: Tensor,
                  H: Tensor,
                  R: Tensor,
                  cholesky=True,
                  ) -> Tuple:
    """
    :param y: measurement, (ydim,)
    :param yhat: predicted measurement, (ydim,)
    :param xhat: predicted state, (xdim,)
    :param Vhat: predicted covariance, (xdim, xdim)
    :param H: measurement matrix, (ydim, xdim)
    :param R: measurement noise covariance, (ydim, ydim)
    :param cholesky: True if Vhat is Cholesky form, default=True
    :return:
        x: posterior mean, (xdim, batch)
        V: posterior covariance or its Cholesky, (xdim, xdim)
    """
    e = y - yhat
    if cholesky:
        Lhat = Vhat
        Vhat = Lhat.mm(Lhat.t())
    else:
        Lhat = linalg.cholesky(Vhat)
    HL = H.mm(Lhat)
    S = HL.mm(HL.t()) + R  # HVH' + R

    L = linalg.cholesky(S)
    G = H.mm(Vhat).cholesky_solve(L).t()  # L^{-1}HV, gain K = VH'S^{-1} = VH'(LL')^{-1} = VH'L'^{-1}L^{-1} = G L^{-1}
    x = xhat + G.mm(e.cholesky_solve(L))
    # V = (I - KH) Vhat (I - KH)' + K R K'
    eye = torch.eye(Vhat.shape[0])
    IminusKH = eye - G.mm(H.cholesky_solve(L))
    IminusKHLhat = IminusKH.mm(Lhat)
    KR = G.mm(R.sqrt().cholesky_solve(L))  # R's supposed to be diagonal
    V = IminusKHLhat.mm(IminusKHLhat.t()) + KR.mm(KR.t())
    if cholesky:
        V = linalg.cholesky(V)
        # print(torch.linalg.eigvalsh(V))

    return x, V
//  [python]
def forward(self, Xnew, full_cov=False, noiseless=True):
        r"""
        Computes the mean and covariance matrix (or variance) of Gaussian Process
        posterior on a test input data :math:`X_{new}`:

        .. math:: p(f^* \mid X_{new}, X, y, k, \epsilon) = \mathcal{N}(loc, cov).

        .. note:: The noise parameter ``noise`` (:math:`\epsilon`) together with
            kernel's parameters have been learned from a training procedure (MCMC or
            SVI).

        :param torch.Tensor Xnew: A input data for testing. Note that
            ``Xnew.shape[1:]`` must be the same as ``self.X.shape[1:]``.
        :param bool full_cov: A flag to decide if we want to predict full covariance
            matrix or just variance.
        :param bool noiseless: A flag to decide if we want to include noise in the
            prediction output or notest_cov[i:i+1]t.
        :returns: loc and covariance matrix (or variance) of :math:`p(f^*(X_{new}))`
        :rtype: tuple(torch.Tensor, torch.Tensor)
        """
        self._check_Xnew_shape(Xnew)
        self.set_mode("guide")

        N = self.X.size(0)
        Kff = self.kernel(self.X).contiguous()
        if self.noise.dim() ==1:
            Kff.view(-1)[:: N + 1] += self.jitter + self.noise
        elif self.noise.dim() ==2:
            Kff = Kff + self.noise
            Kff.view(-1)[:: N + 1] += self.jitter 
        # Kff.view(-1)[:: N + 1] += self.jitter + self.noise  # add noise to the diagonal
        Lff = torch.cholesky(Kff)

        y_residual = self.y - self.mean_function(self.X)
        loc, cov = conditional(
            Xnew,
            self.X,
            self.kernel,
            y_residual,
            None,
            Lff,
            full_cov,
            jitter=self.jitter,
        )

        if full_cov and not noiseless:
            M = Xnew.size(0)
            cov = cov.contiguous()
            cov.view(-1, M * M)[:, :: M + 1] += self.noise  # add noise to the diagonal
        if not full_cov and not noiseless:
            cov = cov + self.noise.abs()

        return loc + self.mean_function(Xnew), cov
//  [python]
def forward(self, Xnew, full_cov=False, noiseless=True):
        r"""
        Computes the mean and covariance matrix (or variance) of Gaussian Process
        posterior on a test input data :math:`X_{new}`:

        .. math:: p(f^* \mid X_{new}, X, y, k, \epsilon) = \mathcal{N}(loc, cov).

        .. note:: The noise parameter ``noise`` (:math:`\epsilon`) together with
            kernel's parameters have been learned from a training procedure (MCMC or
            SVI).

        :param torch.Tensor Xnew: A input data for testing. Note that
            ``Xnew.shape[1:]`` must be the same as ``self.X.shape[1:]``.
        :param bool full_cov: A flag to decide if we want to predict full covariance
            matrix or just variance.
        :param bool noiseless: A flag to decide if we want to include noise in the
            prediction output or notest_cov[i:i+1]t.
        :returns: loc and covariance matrix (or variance) of :math:`p(f^*(X_{new}))`
        :rtype: tuple(torch.Tensor, torch.Tensor)
        """
        self._check_Xnew_shape(Xnew)
        self.set_mode("guide")

        N = self.X.size(0)
        Kff = self.kernel(self.X).contiguous()
        if self.noise.dim() ==1:
            Kff.view(-1)[:: N + 1] += self.jitter + self.noise
        elif self.noise.dim() ==2:
            Kff = Kff + self.noise
            Kff.view(-1)[:: N + 1] += self.jitter 
        # Kff.view(-1)[:: N + 1] += self.jitter + self.noise  # add noise to the diagonal
        Lff = torch.linalg.cholesky(Kff)

        y_residual = self.y - self.mean_function(self.X)
        loc, cov = conditional(
            Xnew,
            self.X,
            self.kernel,
            y_residual,
            None,
            Lff,
            full_cov,
            jitter=self.jitter,
        )

        if full_cov and not noiseless:
            M = Xnew.size(0)
            cov = cov.contiguous()
            cov.view(-1, M * M)[:, :: M + 1] += self.noise  # add noise to the diagonal
        if not full_cov and not noiseless:
            cov = cov + self.noise.abs()

        return loc + self.mean_function(Xnew), cov
//  [python]
def model(self):
        self.set_mode("model")
        N = self.X.size(0)
        x_noise = pyro.sample('obs',dist.Normal(torch.zeros(N),self.xerr**0.5).to_event(1))
        if self.X.dim()<=1:
            X_noisy = (self.X+x_noise)
        else:
            X_noisy = self.X
            X_noisy[:,0] += x_noise
        
        Kff = self.kernel(X_noisy)
        Kff.view(-1)[:: N + 1] += self.jitter + self.noise  # add noise to diagonal
        Lff = torch.cholesky(Kff)
        zero_loc = X_noisy.new_zeros(X_noisy.size(0))
        f_loc = zero_loc + self.mean_function(X_noisy)
        if self.y is None:
            f_var = Lff.pow(2).sum(dim=-1)
            return f_loc, f_var
        else:

            return pyro.sample(
                self._pyro_get_fullname("y"),
                dist.MultivariateNormal(f_loc, scale_tril=Lff)
                .expand_by(self.y.shape[:-1])
                .to_event(self.y.dim() - 1),
                obs=self.y,
            )
//  [python]
def model(self):
        self.set_mode("model")
        N = self.X.size(0)
        x_noise = pyro.sample('obs',dist.Normal(torch.zeros(N),self.xerr**0.5).to_event(1))
        if self.X.dim()<=1:
            X_noisy = (self.X+x_noise)
        else:
            X_noisy = self.X
            X_noisy[:,0] += x_noise
        
        Kff = self.kernel(X_noisy)
        Kff.view(-1)[:: N + 1] += self.jitter + self.noise  # add noise to diagonal
        Lff = torch.linalg.cholesky(Kff)
        zero_loc = X_noisy.new_zeros(X_noisy.size(0))
        f_loc = zero_loc + self.mean_function(X_noisy)
        if self.y is None:
            f_var = Lff.pow(2).sum(dim=-1)
            return f_loc, f_var
        else:

            return pyro.sample(
                self._pyro_get_fullname("y"),
                dist.MultivariateNormal(f_loc, scale_tril=Lff)
                .expand_by(self.y.shape[:-1])
                .to_event(self.y.dim() - 1),
                obs=self.y,
            )
//  [python]
def covs_to_prec(covs):
    cov_chol = torch.cholesky(covs)
    Id = torch.eye(covs.shape[-1], dtype=covs.dtype, device=covs.device)
    prec_chol = torch.triangular_solve(Id, cov_chol, upper=False)[0]
    return prec_chol.transpose(-1, -2)
//  [python]
def covs_to_prec(covs):
    cov_chol = torch.linalg.cholesky(covs)
    Id = torch.eye(covs.shape[-1], dtype=covs.dtype, device=covs.device)
    prec_chol = torch.linalg.solve_triangular(cov_chol, Id, upper=False)
    return prec_chol.transpose(-1, -2)
//  [python]
def iter_sample(self, noiseless=True):
        r"""
        Iteratively constructs a sample from the Gaussian Process posterior.

        Recall that at test input points :math:`X_{new}`, the posterior is
        multivariate Gaussian distributed with mean and covariance matrix
        given by :func:`forward`.

        This method samples lazily from this multivariate Gaussian. The advantage
        of this approach is that later query points can depend upon earlier ones.
        Particularly useful when the querying is to be done by an optimisation
        routine.

        .. note:: The noise parameter ``noise`` (:math:`\epsilon`) together with
            kernel's parameters have been learned from a training procedure (MCMC or
            SVI).

        :param bool noiseless: A flag to decide if we want to add sampling noise
            to the samples beyond the noise inherent in the GP posterior.
        :returns: sampler
        :rtype: function
        """
        noise = self.noise.detach()
        X = self.X.clone().detach()
        y = self.y.clone().detach()
        N = X.size(0)
        Kff = self.kernel(X).contiguous()
        Kff.view(-1)[:: N + 1] += noise  # add noise to the diagonal

        outside_vars = {"X": X, "y": y, "N": N, "Kff": Kff}

        def sample_next(xnew, outside_vars):
            """Repeatedly samples from the Gaussian process posterior,
            conditioning on previously sampled values.
            """
            warn_if_nan(xnew)

            # Variables from outer scope
            X, y, Kff = outside_vars["X"], outside_vars["y"], outside_vars["Kff"]

            # Compute Cholesky decomposition of kernel matrix
            Lff = torch.cholesky(Kff)
            y_residual = y - self.mean_function(X)

            # Compute conditional mean and variance
            loc, cov = conditional(
                xnew, X, self.kernel, y_residual, None, Lff, False, jitter=self.jitter
            )
            if not noiseless:
                cov = cov + noise

            ynew = torchdist.Normal(
                loc + self.mean_function(xnew), cov.sqrt()
            ).rsample()

            # Update kernel matrix
            N = outside_vars["N"]
            Kffnew = Kff.new_empty(N + 1, N + 1)
            Kffnew[:N, :N] = Kff
            cross = self.kernel(X, xnew).squeeze()
            end = self.kernel(xnew, xnew).squeeze()
            Kffnew[N, :N] = cross
            Kffnew[:N, N] = cross
            # No noise, just jitter for numerical stability
            Kffnew[N, N] = end + self.jitter
            # Heuristic to avoid adding degenerate points
            if Kffnew.logdet() > -15.0:
                outside_vars["Kff"] = Kffnew
                outside_vars["N"] += 1
                outside_vars["X"] = torch.cat((X, xnew))
                outside_vars["y"] = torch.cat((y, ynew))

            return ynew

        return lambda xnew: sample_next(xnew, outside_vars)