jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def model(self):
        self.set_mode("model")

        # W = (inv(Luu) @ Kuf).T
        # Qff = Kfu @ inv(Kuu) @ Kuf = W @ W.T
        # Fomulas for each approximation method are
        # DTC:  y_cov = Qff + noise,                   trace_term = 0
        # FITC: y_cov = Qff + diag(Kff - Qff) + noise, trace_term = 0
        # VFE:  y_cov = Qff + noise,                   trace_term = tr(Kff-Qff) / noise
        # y_cov = W @ W.T + D
        # trace_term is added into log_prob

        N = self.X.size(0)
        M = self.Xu.size(0)
        Kuu = self.kernel(self.Xu).contiguous()
        Kuu.view(-1)[:: M + 1] += self.jitter  # add jitter to the diagonal
        Luu = torch.linalg.cholesky(Kuu)
        Kuf = self.kernel(self.Xu, self.X)
        W = torch.linalg.solve_triangular(Luu, Kuf, upper=False).t()

        D = self.noise.expand(N)
        if self.approx == "FITC" or self.approx == "VFE":
            Kffdiag = self.kernel(self.X, diag=True)
            Qffdiag = W.pow(2).sum(dim=-1)
            if self.approx == "FITC":
                D = D + Kffdiag - Qffdiag
            else:  # approx = "VFE"
                trace_term = (Kffdiag - Qffdiag).sum() / self.noise
                trace_term = trace_term.clamp(min=0)

        zero_loc = self.X.new_zeros(N)
        f_loc = zero_loc + self.mean_function(self.X)
        if self.y is None:
            f_var = D + W.pow(2).sum(dim=-1)
            return f_loc, f_var
        else:
            if self.approx == "VFE":
                pyro.factor(self._pyro_get_fullname("trace_term"), -trace_term / 2.0)

            return pyro.sample(
                self._pyro_get_fullname("y"),
                dist.LowRankMultivariateNormal(f_loc, W, D)
                .expand_by(self.y.shape[:-1])
                .to_event(self.y.dim() - 1),
                obs=self.y,
            )
//  [python]
def _solve_linear_system(self, A, B):
        """Solves linear system AX = B.
        If B is a tuple (B1, B2, ...), returns tuple (X1, X2, ...).
        Otherwise returns X.
        """
        B_sizes = None
        # If B is a tuple, concatenate into single tensor:
        if isinstance(B, (tuple, list)):
            B_sizes = list(map(lambda x: x.size(-1), B))
            B = torch.cat(B, dim=-1)
        # Ensure B is 2D (bxmxn):
        if len(B.size()) == 2:
            B = B.unsqueeze(-1)
        try:  # Batchwise Cholesky solve
            A_decomp = torch.cholesky(A, upper=False)
            X = torch.cholesky_solve(B, A_decomp, upper=False)  # bxmxn
        except:  # Revert to loop if batchwise solve fails
            X = torch.zeros_like(B)
            for i in range(A.size(0)):
                try:  # Cholesky solve
                    A_decomp = torch.cholesky(A[i, ...], upper=False)
                    X[i, ...] = torch.cholesky_solve(B[i, ...], A_decomp,
                                                     upper=False)  # mxn
                except:  # Revert to LU solve
                    # X[i, ...], _ = torch.solve(B[i, ...], A[i, ...])  # mxn
                    X[i, ...] = torch.linalg.solve(A[i, ...], B[i, ...])
        if B_sizes is not None:
            X = X.split(B_sizes, dim=-1)
        return X
//  [python]
def _solve_linear_system(self, A, B):
        """Solves linear system AX = B.
        If B is a tuple (B1, B2, ...), returns tuple (X1, X2, ...).
        Otherwise returns X.
        """
        B_sizes = None
        # If B is a tuple, concatenate into single tensor:
        if isinstance(B, (tuple, list)):
            B_sizes = list(map(lambda x: x.size(-1), B))
            B = torch.cat(B, dim=-1)
        # Ensure B is 2D (bxmxn):
        if len(B.size()) == 2:
            B = B.unsqueeze(-1)
        try:  # Batchwise Cholesky solve
            A_decomp = torch.linalg.cholesky(A, upper=False)
            X = torch.cholesky_solve(B, A_decomp, upper=False)  # bxmxn
        except:  # Revert to loop if batchwise solve fails
            X = torch.zeros_like(B)
            for i in range(A.size(0)):
                try:  # Cholesky solve
                    A_decomp = torch.linalg.cholesky(A[i, ...], upper=False)
                    X[i, ...] = torch.cholesky_solve(B[i, ...], A_decomp,
                                                     upper=False)  # mxn
                except:  # Revert to LU solve
                    # X[i, ...], _ = torch.solve(B[i, ...], A[i, ...])  # mxn
                    X[i, ...] = torch.linalg.solve(A[i, ...], B[i, ...])
        if B_sizes is not None:
            X = X.split(B_sizes, dim=-1)
        return X
//  [python]
def test_natgrad(self, D=5):
        mu = torch.randn(D)
        cov = torch.randn(D, D).tril_()
        dist = MultivariateNormal(mu, CholLinearOperator(TriangularLinearOperator(cov)))
        sample = dist.sample()

        v_dist = NaturalVariationalDistribution(D)
        v_dist.initialize_variational_distribution(dist)
        mu = v_dist().mean.detach()

        v_dist().log_prob(sample).squeeze().backward()

        eta1 = mu.clone().requires_grad_(True)
        eta2 = (mu[:, None] * mu + cov @ cov.t()).requires_grad_(True)
        L = torch.cholesky(eta2 - eta1[:, None] * eta1)
        dist2 = MultivariateNormal(eta1, CholLinearOperator(TriangularLinearOperator(L)))
        dist2.log_prob(sample).squeeze().backward()

        assert torch.allclose(v_dist.natural_vec.grad, eta1.grad)
        assert torch.allclose(v_dist.natural_mat.grad, eta2.grad)
//  [python]
def test_natgrad(self, D=5):
        mu = torch.randn(D)
        cov = torch.randn(D, D).tril_()
        dist = MultivariateNormal(mu, CholLinearOperator(TriangularLinearOperator(cov)))
        sample = dist.sample()

        v_dist = NaturalVariationalDistribution(D)
        v_dist.initialize_variational_distribution(dist)
        mu = v_dist().mean.detach()

        v_dist().log_prob(sample).squeeze().backward()

        eta1 = mu.clone().requires_grad_(True)
        eta2 = (mu[:, None] * mu + cov @ cov.t()).requires_grad_(True)
        L = torch.linalg.cholesky(eta2 - eta1[:, None] * eta1)
        dist2 = MultivariateNormal(eta1, CholLinearOperator(TriangularLinearOperator(L)))
        dist2.log_prob(sample).squeeze().backward()

        assert torch.allclose(v_dist.natural_vec.grad, eta1.grad)
        assert torch.allclose(v_dist.natural_mat.grad, eta2.grad)
//  [python]
def solve_dual(self, P, c):
        m = P.shape[1]  # number of columns = batch size

        # construct lhs matrix P* P + m I
        lhs_mat = torch.mm(P.t(), P)
        lhs_mat.diagonal().add_(m)

        # solve positive-definite linear system using Cholesky factorization
        lhs_factor = torch.cholesky(lhs_mat)
        rhs_col = c.unsqueeze(1)  # make rhs a column vector, so that cholesky_solve works
        return torch.cholesky_solve(rhs_col, lhs_factor)
//  [python]
def solve_dual(self, P, c):
        m = P.shape[1]  # number of columns = batch size

        # construct lhs matrix P* P + m I
        lhs_mat = torch.mm(P.t(), P)
        lhs_mat.diagonal().add_(m)

        # solve positive-definite linear system using Cholesky factorization
        lhs_factor = torch.linalg.cholesky(lhs_mat)
        rhs_col = c.unsqueeze(1)  # make rhs a column vector, so that cholesky_solve works
        return torch.cholesky_solve(rhs_col, lhs_factor)
//  [python]
def test_sde_time_diffusion_score(Diffusion, predict_gaussian_noise):
    np.random.seed(132)
    torch.manual_seed(123)

    N = 10_000
    T = 10
    t1 = 1

    x = torch.randn(1, T, 1).repeat(N, 1, 1).square().to(device) + 1
    t = torch.linspace(0, 1, T).view(1, -1, 1).repeat(N, 1, 1).to(device)

    class SDE(torch.nn.Module):
        sde_type = 'ito'
        noise_type = 'general'

        def __init__(self, beta_fn, cov):
            super().__init__()
            self.beta_fn = beta_fn
            self.L = torch.cholesky(cov)

        def f(self, t, x):
            return -0.5 * self.beta_fn(t) * x

        def g(self, t, x):
            return self.beta_fn(t).sqrt() * self.L


    for ratio in [0.01, 0.3, 0.9]:
        i = torch.Tensor([ratio * t1]).view(1, 1, 1).repeat(N, T, 1).to(device)

        beta_fn = BetaLinear(0.1, 10)
        sdediff = Diffusion(
            dim=1,
            t1=t1,
            beta_fn=beta_fn,
            predict_gaussian_noise=predict_gaussian_noise,
        ).to(device)

        y, diff_noise, diff_mean, diff_std, diff_cov = sdediff(x, t=t, i=i, _return_all=True)
        L = torch.cholesky(diff_cov[0])

        y = y.requires_grad_(True)

        assert len(diff_mean.unique()) == 10
        assert torch.all(diff_cov[0] == diff_cov[1])

        diff_mean = diff_mean[0,:,0]
        diff_cov_tilde = (diff_cov * diff_std**2)[0]

        def model(*args, noise=None, **kwargs):
            # "Fake" model that perfectly predicts the noise
            # In case `predict_gaussian_noise=True`, undo covariance
            if predict_gaussian_noise:
                noise = torch.inverse(L) @ noise
            return noise

        model_score = sdediff._get_score(model, x, i=i, t=t, L=L, noise=diff_noise)

        # Statistics from diffusion function vs. empirical covariance of diffused values
        empirical_diff_mean = y.mean([0, 2])
        empirical_diff_cov = torch.cov(y.squeeze(-1).T)
        assert torch.allclose(empirical_diff_mean, diff_mean, atol=0.05)
        assert torch.allclose(empirical_diff_cov, diff_cov_tilde, atol=0.05)

        # Empirical score vs. score from diffusion function
        diff_dist = torch.distributions.MultivariateNormal(diff_mean, diff_cov_tilde)
        log_prob = diff_dist.log_prob(y.squeeze(-1))
        empirical_diff_score = torch.autograd.grad(log_prob.sum(), y)[0]
        assert torch.allclose(empirical_diff_score, model_score, atol=0.05, rtol=0.05)

        # True theoretical covariance vs. calculated in diffusion function
        time_cov = sdediff.noise.covariance(t)
        assert torch.allclose(time_cov[0], diff_cov, atol=0.05)
        true_cov = (1 - torch.exp(-beta_fn.integral(i[0][0]))) * time_cov[0]
        assert torch.allclose(true_cov, diff_cov_tilde, atol=0.05)

        # SDE sequential computation vs. direct covariance computation
        sde = SDE(beta_fn=beta_fn, cov=time_cov)
        times = torch.Tensor([0, ratio * t1]).to(y)
        with torch.no_grad():
            true_y = torchsde.sdeint(sde, x.squeeze(-1), times, dt=5e-4)[-1]
        sde_cov = torch.cov(true_y.T)
        assert torch.allclose(true_cov, sde_cov, atol=0.05)
//  [python]
def test_natgrad(self, D=5):
        mu = torch.randn(D)
        cov = torch.randn(D, D)
        cov = cov @ cov.t()
        dist = MultivariateNormal(mu, CholLinearOperator(TriangularLinearOperator(torch.cholesky(cov))))
        sample = dist.sample()

        v_dist = TrilNaturalVariationalDistribution(D, mean_init_std=0.0)
        v_dist.initialize_variational_distribution(dist)
        v_dist().log_prob(sample).squeeze().backward()
        dout_dnat1 = v_dist.natural_vec.grad
        dout_dnat2 = v_dist.natural_tril_mat.grad

        # mean_init_std=0. because we need to ensure both have the same distribution
        v_dist_ref = NaturalVariationalDistribution(D, mean_init_std=0.0)
        v_dist_ref.initialize_variational_distribution(dist)
        v_dist_ref().log_prob(sample).squeeze().backward()
        dout_dnat1_noforward_ref = v_dist_ref.natural_vec.grad
        dout_dnat2_noforward_ref = v_dist_ref.natural_mat.grad

        def f(natural_vec, natural_tril_mat):
            "Transform natural_tril_mat to L"
            Sigma = torch.inverse(-2 * natural_tril_mat)
            mu = natural_vec
            return mu, torch.cholesky(Sigma).inverse().tril()

        (mu_ref, natural_tril_mat_ref), (dout_dmu_ref, dout_dnat2_ref) = jvp(
            f,
            (v_dist_ref.natural_vec.detach(), v_dist_ref.natural_mat.detach()),
            (dout_dnat1_noforward_ref, dout_dnat2_noforward_ref),
        )

        assert torch.allclose(natural_tril_mat_ref, v_dist.natural_tril_mat), "Sigma transformation"
        assert torch.allclose(dout_dnat2_ref, dout_dnat2), "Sigma gradient"

        assert torch.allclose(mu_ref, v_dist.natural_vec), "mu transformation"
        assert torch.allclose(dout_dmu_ref, dout_dnat1), "mu gradient"
//  [python]
def test_natgrad(self, D=5):
        mu = torch.randn(D)
        cov = torch.randn(D, D)
        cov = cov @ cov.t()
        dist = MultivariateNormal(mu, CholLinearOperator(TriangularLinearOperator(torch.linalg.cholesky(cov))))
        sample = dist.sample()

        v_dist = TrilNaturalVariationalDistribution(D, mean_init_std=0.0)
        v_dist.initialize_variational_distribution(dist)
        v_dist().log_prob(sample).squeeze().backward()
        dout_dnat1 = v_dist.natural_vec.grad
        dout_dnat2 = v_dist.natural_tril_mat.grad

        # mean_init_std=0. because we need to ensure both have the same distribution
        v_dist_ref = NaturalVariationalDistribution(D, mean_init_std=0.0)
        v_dist_ref.initialize_variational_distribution(dist)
        v_dist_ref().log_prob(sample).squeeze().backward()
        dout_dnat1_noforward_ref = v_dist_ref.natural_vec.grad
        dout_dnat2_noforward_ref = v_dist_ref.natural_mat.grad

        def f(natural_vec, natural_tril_mat):
            "Transform natural_tril_mat to L"
            Sigma = torch.inverse(-2 * natural_tril_mat)
            mu = natural_vec
            return mu, torch.linalg.cholesky(Sigma).inverse().tril()

        (mu_ref, natural_tril_mat_ref), (dout_dmu_ref, dout_dnat2_ref) = jvp(
            f,
            (v_dist_ref.natural_vec.detach(), v_dist_ref.natural_mat.detach()),
            (dout_dnat1_noforward_ref, dout_dnat2_noforward_ref),
        )

        assert torch.allclose(natural_tril_mat_ref, v_dist.natural_tril_mat), "Sigma transformation"
        assert torch.allclose(dout_dnat2_ref, dout_dnat2), "Sigma gradient"

        assert torch.allclose(mu_ref, v_dist.natural_vec), "mu transformation"
        assert torch.allclose(dout_dmu_ref, dout_dnat1), "mu gradient"
//  [python]
def compute_energy(self, z, phi=None, mu=None, cov=None, size_average=True):
        if phi is None:
            phi = to_var(self.phi)
        if mu is None:
            mu = to_var(self.mu)
        if cov is None:
            cov = to_var(self.cov)

        k, D, _ = cov.size()

        z_mu = (z.unsqueeze(1) - mu.unsqueeze(0))

        cov_inverse = []
        det_cov = []
        cov_diag = 0
        eps = 1e-12
        for i in range(k):
            # K x D x D
            cov_k = cov[i] + to_var(torch.eye(D) * eps)
            cov_inverse.append(torch.inverse(cov_k).unsqueeze(0))
            det_cov.append(torch.potrf(cov_k * 2 * math.pi).diag().prod().unsqueeze(0))
            cov_diag = cov_diag + torch.sum(1 / cov_k.diag())

        # K x D x D
        cov_inverse = torch.cat(cov_inverse, dim=0)
        # K
        det_cov = torch.cat(det_cov).cuda()

        # N x K
        exp_term_tmp = -0.5 * torch.sum(torch.sum(z_mu.unsqueeze(-1) * cov_inverse.unsqueeze(0), dim=-2) * z_mu, dim=-1)
        # for stability (logsumexp)
        max_val = torch.max((exp_term_tmp).clamp(min=0), dim=1, keepdim=True)[0]
        exp_term = torch.exp(exp_term_tmp - max_val)
        sample_energy = -max_val.squeeze() - torch.log(
            torch.sum(phi.unsqueeze(0) * exp_term / (torch.sqrt(det_cov)).unsqueeze(0), dim=1) + eps)

        if size_average:
            sample_energy = torch.mean(sample_energy)

        return sample_energy, cov_diag
//  [python]
def compute_energy(self, z, phi=None, mu=None, cov=None, size_average=True):
        if phi is None:
            phi = to_var(self.phi)
        if mu is None:
            mu = to_var(self.mu)
        if cov is None:
            cov = to_var(self.cov)

        k, D, _ = cov.size()

        z_mu = (z.unsqueeze(1) - mu.unsqueeze(0))

        cov_inverse = []
        det_cov = []
        cov_diag = 0
        eps = 1e-12
        for i in range(k):
            # K x D x D
            cov_k = cov[i] + to_var(torch.eye(D) * eps)
            cov_inverse.append(torch.inverse(cov_k).unsqueeze(0))
            det_cov.append(torch.linalg.cholesky(cov_k * 2 * math.pi).diag().prod().unsqueeze(0))
            cov_diag = cov_diag + torch.sum(1 / cov_k.diag())

        # K x D x D
        cov_inverse = torch.cat(cov_inverse, dim=0)
        # K
        det_cov = torch.cat(det_cov).cuda()

        # N x K
        exp_term_tmp = -0.5 * torch.sum(torch.sum(z_mu.unsqueeze(-1) * cov_inverse.unsqueeze(0), dim=-2) * z_mu, dim=-1)
        # for stability (logsumexp)
        max_val = torch.max((exp_term_tmp).clamp(min=0), dim=1, keepdim=True)[0]
        exp_term = torch.exp(exp_term_tmp - max_val)
        sample_energy = -max_val.squeeze() - torch.log(
            torch.sum(phi.unsqueeze(0) * exp_term / (torch.sqrt(det_cov)).unsqueeze(0), dim=1) + eps)

        if size_average:
            sample_energy = torch.mean(sample_energy)

        return sample_energy, cov_diag
//  [python]
def lml(XY_train, UV_train, kind, params):
    """
    lml returns the log marginal likelihood of a GP model over given training data
    Args:
        XY_train: np.array of shape [N_train, 2] with X & Y coordinates of observations
        UV_train: np.array of shape [N_train, 2] with U & V flow observations
        kind: "standard" for velocity GP, "helmholtz" for Helmholtz GP
        params: parameter of the kernel, from the @dataclass TwoKernelGPParams
    """
    if kind=="standard":
        kernel = lambda XY_11, XY_22, params : kernel_fcn_uv_twodata(XY_11, XY_22, params)
    elif kind=="helmholtz":
        kernel = lambda XY_11, XY_22, params : kernel_fcn_helm_twodata(XY_11, XY_22, params)

    train_obs = torch.cat((UV_train[:,0], UV_train[:,1]), dim=0)[:,None]
    M = XY_train.shape[0]
    N = XY_train.shape[0]
    XY_11 = torch.tile(XY_train[:, None], [1, N, 1])
    XY_22 = torch.tile(XY_train[None], [M, 1, 1])
    K_trtr = kernel(XY_11, XY_22, params) + torch.eye(2*M) * params.obs_noise**2
    chol_factor = torch.cholesky(K_trtr)
    logdet = - torch.sum(torch.log(torch.diag(chol_factor)))
    quadratic_term = -0.5 * torch.sum(torch.square(torch.triangular_solve(train_obs, chol_factor, upper=False)[0]))
    ll = logdet + quadratic_term
    return ll
//  [python]
def lml(XY_train, UV_train, kind, params):
    """
    lml returns the log marginal likelihood of a GP model over given training data
    Args:
        XY_train: np.array of shape [N_train, 2] with X & Y coordinates of observations
        UV_train: np.array of shape [N_train, 2] with U & V flow observations
        kind: "standard" for velocity GP, "helmholtz" for Helmholtz GP
        params: parameter of the kernel, from the @dataclass TwoKernelGPParams
    """
    if kind=="standard":
        kernel = lambda XY_11, XY_22, params : kernel_fcn_uv_twodata(XY_11, XY_22, params)
    elif kind=="helmholtz":
        kernel = lambda XY_11, XY_22, params : kernel_fcn_helm_twodata(XY_11, XY_22, params)

    train_obs = torch.cat((UV_train[:,0], UV_train[:,1]), dim=0)[:,None]
    M = XY_train.shape[0]
    N = XY_train.shape[0]
    XY_11 = torch.tile(XY_train[:, None], [1, N, 1])
    XY_22 = torch.tile(XY_train[None], [M, 1, 1])
    K_trtr = kernel(XY_11, XY_22, params) + torch.eye(2*M) * params.obs_noise**2
    chol_factor = torch.linalg.cholesky(K_trtr)
    logdet = - torch.sum(torch.log(torch.linalg.diagonal(chol_factor)))
    quadratic_term = -0.5 * torch.sum(torch.square(torch.linalg.solve_triangular(chol_factor, train_obs, upper=False)))
    ll = logdet + quadratic_term
    return ll
//  [python]
def __init__(self, *, representation: LBEigenspaceWithBasis):
        super().__init__()
        self.representation = representation
        dim = self.representation.dimension
        self.character = self.representation.phase_function
        gram_is_spd = False
        attempts = 0
        while not gram_is_spd:
            attempts += 1
            self.translations = self.representation.manifold.rand(dim**2)
            ratios = self.representation.manifold.pairwise_embed(self.translations, self.translations)
            gram = self.character.forward(ratios).reshape(dim**2, dim**2)
            try:
                self.coeffs = torch.inverse(torch.cholesky(gram))
                gram_is_spd = True
            except RuntimeError:
                if attempts >= 3:
                    raise
//  [python]
def __init__(self, *, representation: LBEigenspaceWithBasis):
        super().__init__()
        self.representation = representation
        dim = self.representation.dimension
        self.character = self.representation.phase_function
        gram_is_spd = False
        attempts = 0
        while not gram_is_spd:
            attempts += 1
            self.translations = self.representation.manifold.rand(dim**2)
            ratios = self.representation.manifold.pairwise_embed(self.translations, self.translations)
            gram = self.character.forward(ratios).reshape(dim**2, dim**2)
            try:
                self.coeffs = torch.linalg.inv(torch.linalg.cholesky(gram))
                gram_is_spd = True
            except RuntimeError:
                if attempts >= 3:
                    raise
//  [python]
def posterior_divergence_forloop(XY_test, XY_train, UV_train, kernel, params):
    """
    posterior_divergence_forloop computes the divergence posterior mean and variance at 
    test points for either the velocity GP (kernel='standard') or Helmholtz GP (kernel='helmholtz')
    Args:
        XY_test: torch.tensor of shape [N_test, 2]
        XY_train: torch.tensor of shape [N_train, 2]
        UV_train: torch.tensor of shape [N_train, 2]
        kernel: "helmholtz" or "standard"
        params: parameter of the kernel, from the @dataclass TwoKernelGPParams
    Returns:   
        mean_div: torch.tensor of shape [N_test, 1]
        var_div: torch.tensor of shape [N_test, 1]
    """
    
    mean_div = []
    var_div = []

    # Compute cholesky factor of train matrix 
    M = XY_train.shape[0]
    N = XY_train.shape[0]
    XY_11 = torch.tile(XY_train[:, None], [1, N, 1])
    XY_22 = torch.tile(XY_train[None], [M, 1, 1])
    kernel_fn = kernel_fcn_uv_twodata if kernel=="standard" else kernel_fcn_helm_twodata
    prior_obs = kernel_fn(XY_11, XY_22, params) + torch.eye(2*M)*params.obs_noise**2
    train_chol = torch.cholesky(prior_obs)
    
    for i in range(XY_test.shape[0]):
        XY_test_i = XY_test[i:i+1,:].clone().requires_grad_(True)
        XY_test_j = XY_test[i:i+1,:].clone().requires_grad_(True)
        # mean_uv_i, cov_uv_i, _ = posterior_kernel_twodata(XY_test_i, XY_test_j, UV_train, XY_train, kernel, params)
        mean_uv_i, cov_uv_i = posterior_kernel_twodata_cached(XY_test_i, XY_test_j, UV_train, XY_train, train_chol, kernel, params) # shapes [2] and [2,2]
        
        mean_div_u_i = torch.autograd.grad(mean_uv_i[0], XY_test_i, create_graph=True)[0][:,0]
        mean_div_v_i = torch.autograd.grad(mean_uv_i[1], XY_test_i, create_graph=True)[0][:,1]
        mean_div_i = mean_div_u_i + mean_div_v_i
        mean_div.append(mean_div_i)

        var_div_uu_i = torch.autograd.grad(torch.autograd.grad(cov_uv_i[0,0], XY_test_i, create_graph=True)[0][:,0], XY_test_j, create_graph=True)[0][:,0]
        var_div_uv_i = torch.autograd.grad(torch.autograd.grad(cov_uv_i[0,1], XY_test_i, create_graph=True)[0][:,0], XY_test_j, create_graph=True)[0][:,1]
        var_div_vu_i = torch.autograd.grad(torch.autograd.grad(cov_uv_i[1,0], XY_test_i, create_graph=True)[0][:,1], XY_test_j, create_graph=True)[0][:,0]
        var_div_vv_i = torch.autograd.grad(torch.autograd.grad(cov_uv_i[1,1], XY_test_i, create_graph=True)[0][:,1], XY_test_j, create_graph=True)[0][:,1]
        
        var_div_i = var_div_uu_i + var_div_uv_i + var_div_vu_i + var_div_vv_i
        var_div.append(var_div_i)
        
    mean_div = torch.tensor(mean_div)[:,None]
    var_div = torch.tensor(var_div)[:,None]

    return mean_div, var_div
//  [python]
def posterior_divergence_forloop(XY_test, XY_train, UV_train, kernel, params):
    """
    posterior_divergence_forloop computes the divergence posterior mean and variance at 
    test points for either the velocity GP (kernel='standard') or Helmholtz GP (kernel='helmholtz')
    Args:
        XY_test: torch.tensor of shape [N_test, 2]
        XY_train: torch.tensor of shape [N_train, 2]
        UV_train: torch.tensor of shape [N_train, 2]
        kernel: "helmholtz" or "standard"
        params: parameter of the kernel, from the @dataclass TwoKernelGPParams
    Returns:   
        mean_div: torch.tensor of shape [N_test, 1]
        var_div: torch.tensor of shape [N_test, 1]
    """
    
    mean_div = []
    var_div = []

    # Compute cholesky factor of train matrix 
    M = XY_train.shape[0]
    N = XY_train.shape[0]
    XY_11 = torch.tile(XY_train[:, None], [1, N, 1])
    XY_22 = torch.tile(XY_train[None], [M, 1, 1])
    kernel_fn = kernel_fcn_uv_twodata if kernel=="standard" else kernel_fcn_helm_twodata
    prior_obs = kernel_fn(XY_11, XY_22, params) + torch.eye(2*M)*params.obs_noise**2
    train_chol = torch.linalg.cholesky(prior_obs)
    
    for i in range(XY_test.shape[0]):
        XY_test_i = XY_test[i:i+1,:].clone().requires_grad_(True)
        XY_test_j = XY_test[i:i+1,:].clone().requires_grad_(True)
        # mean_uv_i, cov_uv_i, _ = posterior_kernel_twodata(XY_test_i, XY_test_j, UV_train, XY_train, kernel, params)
        mean_uv_i, cov_uv_i = posterior_kernel_twodata_cached(XY_test_i, XY_test_j, UV_train, XY_train, train_chol, kernel, params) # shapes [2] and [2,2]
        
        mean_div_u_i = torch.autograd.grad(mean_uv_i[0], XY_test_i, create_graph=True)[0][:,0]
        mean_div_v_i = torch.autograd.grad(mean_uv_i[1], XY_test_i, create_graph=True)[0][:,1]
        mean_div_i = mean_div_u_i + mean_div_v_i
        mean_div.append(mean_div_i)

        var_div_uu_i = torch.autograd.grad(torch.autograd.grad(cov_uv_i[0,0], XY_test_i, create_graph=True)[0][:,0], XY_test_j, create_graph=True)[0][:,0]
        var_div_uv_i = torch.autograd.grad(torch.autograd.grad(cov_uv_i[0,1], XY_test_i, create_graph=True)[0][:,0], XY_test_j, create_graph=True)[0][:,1]
        var_div_vu_i = torch.autograd.grad(torch.autograd.grad(cov_uv_i[1,0], XY_test_i, create_graph=True)[0][:,1], XY_test_j, create_graph=True)[0][:,0]
        var_div_vv_i = torch.autograd.grad(torch.autograd.grad(cov_uv_i[1,1], XY_test_i, create_graph=True)[0][:,1], XY_test_j, create_graph=True)[0][:,1]
        
        var_div_i = var_div_uu_i + var_div_uv_i + var_div_vu_i + var_div_vv_i
        var_div.append(var_div_i)
        
    mean_div = torch.tensor(mean_div)[:,None]
    var_div = torch.tensor(var_div)[:,None]

    return mean_div, var_div
//  [python]
def finalize(self, loss, target_labels):
        """
        Compute the Hessian of the parameters wrt ``loss``

        :param torch.Tensor loss: the output of evaluating a loss function such as
                                  `pyro.infer.Trace_ELBO().differentiable_loss` on the model, guide and design.
        :param list target_labels: list indicating the sample sites that are targets, i.e. for which information gain
                                   should be measured.
        """
        self.eval()
        for l, mu_l in self.means.items():
            if l not in target_labels:
                continue
            hess_l = self._hessian_diag(loss, mu_l, event_shape=(self.w_sizes[l],))
            cov_l = rinverse(hess_l)
            self.scale_trils[l] = torch.cholesky(cov_l.transpose(-2, -1)).transpose(-2, -1)
//  [python]
def finalize(self, loss, target_labels):
        """
        Compute the Hessian of the parameters wrt ``loss``

        :param torch.Tensor loss: the output of evaluating a loss function such as
                                  `pyro.infer.Trace_ELBO().differentiable_loss` on the model, guide and design.
        :param list target_labels: list indicating the sample sites that are targets, i.e. for which information gain
                                   should be measured.
        """
        self.eval()
        for l, mu_l in self.means.items():
            if l not in target_labels:
                continue
            hess_l = self._hessian_diag(loss, mu_l, event_shape=(self.w_sizes[l],))
            cov_l = rinverse(hess_l)
            self.scale_trils[l] = torch.linalg.cholesky(cov_l.transpose(-2, -1)).transpose(-2, -1)
//  [python]
def posterior_vorticity_forloop(XY_test, XY_train, UV_train, kernel, params):
    """
    posterior_vorticity_forloop computes the vorticity posterior mean and variance at 
    test points for either the velocity GP (kernel='standard') or Helmholtz GP (kernel='helmholtz')
    Args:
        XY_test: torch.tensor of shape [N_test, 2]
        XY_train: torch.tensor of shape [N_train, 2]
        UV_train: torch.tensor of shape [N_train, 2]
        kernel: "helmholtz" or "standard"
        params: parameter of the kernel, from the @dataclass TwoKernelGPParams
    Returns:   
        mean_vort: torch.tensor of shape [N_test, 1]
        var_vort: torch.tensor of shape [N_test, 1]
    """
    
    mean_vort = []
    var_vort = []

    # Compute cholesky factor of train matrix 
    M = XY_train.shape[0]
    N = XY_train.shape[0]
    XY_11 = torch.tile(XY_train[:, None], [1, N, 1])
    XY_22 = torch.tile(XY_train[None], [M, 1, 1])
    kernel_fn = kernel_fcn_uv_twodata if kernel=="standard" else kernel_fcn_helm_twodata
    prior_obs = kernel_fn(XY_11, XY_22, params) + torch.eye(2*M)*params.obs_noise**2
    train_chol = torch.cholesky(prior_obs)
    for i in range(XY_test.shape[0]):
        XY_test_i = XY_test[i:i+1,:].clone().requires_grad_(True)
        XY_test_j = XY_test[i:i+1,:].clone().requires_grad_(True)
        
        mean_uv_i, cov_uv_i = posterior_kernel_twodata_cached(XY_test_i, XY_test_j, UV_train, XY_train, train_chol, kernel, params) # shapes [2] and [2,2]
        
        mean_vort_u_i = torch.autograd.grad(mean_uv_i[0], XY_test_i, create_graph=True)[0][:,1]
        mean_vort_v_i = -torch.autograd.grad(mean_uv_i[1], XY_test_i, create_graph=True)[0][:,0]
        mean_vort_i = mean_vort_u_i + mean_vort_v_i
        mean_vort.append(mean_vort_i)

        var_vort_uu_i = torch.autograd.grad(torch.autograd.grad(cov_uv_i[0,0], XY_test_i, create_graph=True)[0][:,1], XY_test_j, create_graph=True)[0][:,1]
        var_vort_uv_i = -torch.autograd.grad(torch.autograd.grad(cov_uv_i[0,1], XY_test_i, create_graph=True)[0][:,1], XY_test_j, create_graph=True)[0][:,0]
        var_vort_vu_i = -torch.autograd.grad(torch.autograd.grad(cov_uv_i[1,0], XY_test_i, create_graph=True)[0][:,0], XY_test_j, create_graph=True)[0][:,1]
        var_vort_vv_i = torch.autograd.grad(torch.autograd.grad(cov_uv_i[1,1], XY_test_i, create_graph=True)[0][:,0], XY_test_j, create_graph=True)[0][:,0]
        
        var_vort_i = var_vort_uu_i + var_vort_uv_i + var_vort_vu_i + var_vort_vv_i
        var_vort.append(var_vort_i)
        
    mean_vort = torch.tensor(mean_vort)[:,None]
    var_vort = torch.tensor(var_vort)[:,None]
    return mean_vort, var_vort
//  [python]
def posterior_vorticity_forloop(XY_test, XY_train, UV_train, kernel, params):
    """
    posterior_vorticity_forloop computes the vorticity posterior mean and variance at 
    test points for either the velocity GP (kernel='standard') or Helmholtz GP (kernel='helmholtz')
    Args:
        XY_test: torch.tensor of shape [N_test, 2]
        XY_train: torch.tensor of shape [N_train, 2]
        UV_train: torch.tensor of shape [N_train, 2]
        kernel: "helmholtz" or "standard"
        params: parameter of the kernel, from the @dataclass TwoKernelGPParams
    Returns:   
        mean_vort: torch.tensor of shape [N_test, 1]
        var_vort: torch.tensor of shape [N_test, 1]
    """
    
    mean_vort = []
    var_vort = []

    # Compute cholesky factor of train matrix 
    M = XY_train.shape[0]
    N = XY_train.shape[0]
    XY_11 = torch.tile(XY_train[:, None], [1, N, 1])
    XY_22 = torch.tile(XY_train[None], [M, 1, 1])
    kernel_fn = kernel_fcn_uv_twodata if kernel=="standard" else kernel_fcn_helm_twodata
    prior_obs = kernel_fn(XY_11, XY_22, params) + torch.eye(2*M)*params.obs_noise**2
    train_chol = torch.linalg.cholesky(prior_obs)
    for i in range(XY_test.shape[0]):
        XY_test_i = XY_test[i:i+1,:].clone().requires_grad_(True)
        XY_test_j = XY_test[i:i+1,:].clone().requires_grad_(True)
        
        mean_uv_i, cov_uv_i = posterior_kernel_twodata_cached(XY_test_i, XY_test_j, UV_train, XY_train, train_chol, kernel, params) # shapes [2] and [2,2]
        
        mean_vort_u_i = torch.autograd.grad(mean_uv_i[0], XY_test_i, create_graph=True)[0][:,1]
        mean_vort_v_i = -torch.autograd.grad(mean_uv_i[1], XY_test_i, create_graph=True)[0][:,0]
        mean_vort_i = mean_vort_u_i + mean_vort_v_i
        mean_vort.append(mean_vort_i)

        var_vort_uu_i = torch.autograd.grad(torch.autograd.grad(cov_uv_i[0,0], XY_test_i, create_graph=True)[0][:,1], XY_test_j, create_graph=True)[0][:,1]
        var_vort_uv_i = -torch.autograd.grad(torch.autograd.grad(cov_uv_i[0,1], XY_test_i, create_graph=True)[0][:,1], XY_test_j, create_graph=True)[0][:,0]
        var_vort_vu_i = -torch.autograd.grad(torch.autograd.grad(cov_uv_i[1,0], XY_test_i, create_graph=True)[0][:,0], XY_test_j, create_graph=True)[0][:,1]
        var_vort_vv_i = torch.autograd.grad(torch.autograd.grad(cov_uv_i[1,1], XY_test_i, create_graph=True)[0][:,0], XY_test_j, create_graph=True)[0][:,0]
        
        var_vort_i = var_vort_uu_i + var_vort_uv_i + var_vort_vu_i + var_vort_vv_i
        var_vort.append(var_vort_i)
        
    mean_vort = torch.tensor(mean_vort)[:,None]
    var_vort = torch.tensor(var_vort)[:,None]
    return mean_vort, var_vort
//  [python]
def posterior_divergence_forloop_diffphi(XY_test, XY_train, UV_train, params):
    """
    posterior_divergence_forloop_diffphi computes the divergence posterior mean and variance for 
    the Helmholtz GP in an optimized way (by exploting the fact that the divergence of the vorticity 
    component is 0 and so does not have to be taken into account)
    Args:
        XY_test: torch.tensor of shape [N_test, 2]
        XY_train: torch.tensor of shape [N_train, 2]
        UV_train: torch.tensor of shape [N_train, 2]
        params: parameter of the kernel, from the @dataclass TwoKernelGPParams
    Returns:   
        mean_div: torch.tensor of shape [N_test, 1]
        var_div: torch.tensor of shape [N_test, 1]
    """
    
    mean_div = []
    var_div = []
    M = XY_train.shape[0]
    N = XY_train.shape[0]
    XY_11 = torch.tile(XY_train[:, None], [1, N, 1])
    XY_22 = torch.tile(XY_train[None], [M, 1, 1])
    prior_obs = kernel_fcn_helm_twodata(XY_11, XY_22, params) + torch.eye(2*M)*params.obs_noise**2
    train_chol = torch.cholesky(prior_obs)
    for i in range(XY_test.shape[0]):
        XY_test_i = XY_test[i:i+1,:].clone().requires_grad_(True)
        XY_test_j = XY_test[i:i+1,:].clone().requires_grad_(True)
        
        mean_uv_i, cov_uv_i = posterior_components_grid_cached(XY_train, UV_train, XY_test_i, XY_test_j, train_chol, params) # shapes [2] and [2,2]
        
        mean_div_u_i = torch.autograd.grad(mean_uv_i[0], XY_test_i, create_graph=True)[0][:,0]
        mean_div_v_i = torch.autograd.grad(mean_uv_i[1], XY_test_i, create_graph=True)[0][:,1]
        mean_div_i = mean_div_u_i + mean_div_v_i
        mean_div.append(mean_div_i)

        var_div_uu_i = torch.autograd.grad(torch.autograd.grad(cov_uv_i[0,0], XY_test_i, create_graph=True)[0][:,0], XY_test_j, create_graph=True)[0][:,0]
        var_div_uv_i = torch.autograd.grad(torch.autograd.grad(cov_uv_i[0,1], XY_test_i, create_graph=True)[0][:,0], XY_test_j, create_graph=True)[0][:,1]
        var_div_vu_i = torch.autograd.grad(torch.autograd.grad(cov_uv_i[1,0], XY_test_i, create_graph=True)[0][:,1], XY_test_j, create_graph=True)[0][:,0]
        var_div_vv_i = torch.autograd.grad(torch.autograd.grad(cov_uv_i[1,1], XY_test_i, create_graph=True)[0][:,1], XY_test_j, create_graph=True)[0][:,1]
        
        var_div_i = var_div_uu_i + var_div_uv_i + var_div_vu_i + var_div_vv_i
        var_div.append(var_div_i)
        
    mean_div = torch.tensor(mean_div)[:,None]
    var_div = torch.tensor(var_div)[:,None]

    return mean_div, var_div
//  [python]
def posterior_divergence_forloop_diffphi(XY_test, XY_train, UV_train, params):
    """
    posterior_divergence_forloop_diffphi computes the divergence posterior mean and variance for 
    the Helmholtz GP in an optimized way (by exploting the fact that the divergence of the vorticity 
    component is 0 and so does not have to be taken into account)
    Args:
        XY_test: torch.tensor of shape [N_test, 2]
        XY_train: torch.tensor of shape [N_train, 2]
        UV_train: torch.tensor of shape [N_train, 2]
        params: parameter of the kernel, from the @dataclass TwoKernelGPParams
    Returns:   
        mean_div: torch.tensor of shape [N_test, 1]
        var_div: torch.tensor of shape [N_test, 1]
    """
    
    mean_div = []
    var_div = []
    M = XY_train.shape[0]
    N = XY_train.shape[0]
    XY_11 = torch.tile(XY_train[:, None], [1, N, 1])
    XY_22 = torch.tile(XY_train[None], [M, 1, 1])
    prior_obs = kernel_fcn_helm_twodata(XY_11, XY_22, params) + torch.eye(2*M)*params.obs_noise**2
    train_chol = torch.linalg.cholesky(prior_obs)
    for i in range(XY_test.shape[0]):
        XY_test_i = XY_test[i:i+1,:].clone().requires_grad_(True)
        XY_test_j = XY_test[i:i+1,:].clone().requires_grad_(True)
        
        mean_uv_i, cov_uv_i = posterior_components_grid_cached(XY_train, UV_train, XY_test_i, XY_test_j, train_chol, params) # shapes [2] and [2,2]
        
        mean_div_u_i = torch.autograd.grad(mean_uv_i[0], XY_test_i, create_graph=True)[0][:,0]
        mean_div_v_i = torch.autograd.grad(mean_uv_i[1], XY_test_i, create_graph=True)[0][:,1]
        mean_div_i = mean_div_u_i + mean_div_v_i
        mean_div.append(mean_div_i)

        var_div_uu_i = torch.autograd.grad(torch.autograd.grad(cov_uv_i[0,0], XY_test_i, create_graph=True)[0][:,0], XY_test_j, create_graph=True)[0][:,0]
        var_div_uv_i = torch.autograd.grad(torch.autograd.grad(cov_uv_i[0,1], XY_test_i, create_graph=True)[0][:,0], XY_test_j, create_graph=True)[0][:,1]
        var_div_vu_i = torch.autograd.grad(torch.autograd.grad(cov_uv_i[1,0], XY_test_i, create_graph=True)[0][:,1], XY_test_j, create_graph=True)[0][:,0]
        var_div_vv_i = torch.autograd.grad(torch.autograd.grad(cov_uv_i[1,1], XY_test_i, create_graph=True)[0][:,1], XY_test_j, create_graph=True)[0][:,1]
        
        var_div_i = var_div_uu_i + var_div_uv_i + var_div_vu_i + var_div_vv_i
        var_div.append(var_div_i)
        
    mean_div = torch.tensor(mean_div)[:,None]
    var_div = torch.tensor(var_div)[:,None]

    return mean_div, var_div
//  [python]
def posterior_vorticity_forloop_diffpsi(XY_test, XY_train, UV_train, params):
    """
    posterior_divergence_forloop_diffps computes the vorticity posterior mean and variance for 
    the Helmholtz GP in an optimized way (by exploting the fact that the curl of the divergence 
    component is 0 and so does not have to be taken into account)
    Args:
        XY_test: torch.tensor of shape [N_test, 2]
        XY_train: torch.tensor of shape [N_train, 2]
        UV_train: torch.tensor of shape [N_train, 2]
        params: parameter of the kernel, from the @dataclass TwoKernelGPParams
    Returns:   
        mean_vort: torch.tensor of shape [N_test, 1]
        var_vort: torch.tensor of shape [N_test, 1]
    """
    
    mean_vort = []
    var_vort = []
    M = XY_train.shape[0]
    N = XY_train.shape[0]
    XY_11 = torch.tile(XY_train[:, None], [1, N, 1])
    XY_22 = torch.tile(XY_train[None], [M, 1, 1])
    prior_obs = kernel_fcn_helm_twodata(XY_11, XY_22, params) + torch.eye(2*M)*params.obs_noise**2
    train_chol = torch.cholesky(prior_obs)
    for i in range(XY_test.shape[0]):
        XY_test_i = XY_test[i:i+1,:].clone().requires_grad_(True)
        XY_test_j = XY_test[i:i+1,:].clone().requires_grad_(True)
        
        mean_uv_i, cov_uv_i = posterior_components_grid_cached(XY_train, UV_train, XY_test_i, XY_test_j, train_chol, params, curl=True) # shapes [2] and [2,2]
        
        mean_vort_u_i = torch.autograd.grad(mean_uv_i[0], XY_test_i, create_graph=True)[0][:,1]
        mean_vort_v_i = -torch.autograd.grad(mean_uv_i[1], XY_test_i, create_graph=True)[0][:,0]
        mean_vort_i = mean_vort_u_i + mean_vort_v_i
        mean_vort.append(mean_vort_i)

        var_vort_uu_i = torch.autograd.grad(torch.autograd.grad(cov_uv_i[0,0], XY_test_i, create_graph=True)[0][:,1], XY_test_j, create_graph=True)[0][:,1]
        var_vort_uv_i = -torch.autograd.grad(torch.autograd.grad(cov_uv_i[0,1], XY_test_i, create_graph=True)[0][:,1], XY_test_j, create_graph=True)[0][:,0]
        var_vort_vu_i = -torch.autograd.grad(torch.autograd.grad(cov_uv_i[1,0], XY_test_i, create_graph=True)[0][:,0], XY_test_j, create_graph=True)[0][:,1]
        var_vort_vv_i = torch.autograd.grad(torch.autograd.grad(cov_uv_i[1,1], XY_test_i, create_graph=True)[0][:,0], XY_test_j, create_graph=True)[0][:,0]
        
        var_vort_i = var_vort_uu_i + var_vort_uv_i + var_vort_vu_i + var_vort_vv_i
        var_vort.append(var_vort_i)
        
    mean_vort = torch.tensor(mean_vort)[:,None]
    var_vort = torch.tensor(var_vort)[:,None]
    return mean_vort, var_vort
//  [python]
def posterior_vorticity_forloop_diffpsi(XY_test, XY_train, UV_train, params):
    """
    posterior_divergence_forloop_diffps computes the vorticity posterior mean and variance for 
    the Helmholtz GP in an optimized way (by exploting the fact that the curl of the divergence 
    component is 0 and so does not have to be taken into account)
    Args:
        XY_test: torch.tensor of shape [N_test, 2]
        XY_train: torch.tensor of shape [N_train, 2]
        UV_train: torch.tensor of shape [N_train, 2]
        params: parameter of the kernel, from the @dataclass TwoKernelGPParams
    Returns:   
        mean_vort: torch.tensor of shape [N_test, 1]
        var_vort: torch.tensor of shape [N_test, 1]
    """
    
    mean_vort = []
    var_vort = []
    M = XY_train.shape[0]
    N = XY_train.shape[0]
    XY_11 = torch.tile(XY_train[:, None], [1, N, 1])
    XY_22 = torch.tile(XY_train[None], [M, 1, 1])
    prior_obs = kernel_fcn_helm_twodata(XY_11, XY_22, params) + torch.eye(2*M)*params.obs_noise**2
    train_chol = torch.linalg.cholesky(prior_obs)
    for i in range(XY_test.shape[0]):
        XY_test_i = XY_test[i:i+1,:].clone().requires_grad_(True)
        XY_test_j = XY_test[i:i+1,:].clone().requires_grad_(True)
        
        mean_uv_i, cov_uv_i = posterior_components_grid_cached(XY_train, UV_train, XY_test_i, XY_test_j, train_chol, params, curl=True) # shapes [2] and [2,2]
        
        mean_vort_u_i = torch.autograd.grad(mean_uv_i[0], XY_test_i, create_graph=True)[0][:,1]
        mean_vort_v_i = -torch.autograd.grad(mean_uv_i[1], XY_test_i, create_graph=True)[0][:,0]
        mean_vort_i = mean_vort_u_i + mean_vort_v_i
        mean_vort.append(mean_vort_i)

        var_vort_uu_i = torch.autograd.grad(torch.autograd.grad(cov_uv_i[0,0], XY_test_i, create_graph=True)[0][:,1], XY_test_j, create_graph=True)[0][:,1]
        var_vort_uv_i = -torch.autograd.grad(torch.autograd.grad(cov_uv_i[0,1], XY_test_i, create_graph=True)[0][:,1], XY_test_j, create_graph=True)[0][:,0]
        var_vort_vu_i = -torch.autograd.grad(torch.autograd.grad(cov_uv_i[1,0], XY_test_i, create_graph=True)[0][:,0], XY_test_j, create_graph=True)[0][:,1]
        var_vort_vv_i = torch.autograd.grad(torch.autograd.grad(cov_uv_i[1,1], XY_test_i, create_graph=True)[0][:,0], XY_test_j, create_graph=True)[0][:,0]
        
        var_vort_i = var_vort_uu_i + var_vort_uv_i + var_vort_vu_i + var_vort_vv_i
        var_vort.append(var_vort_i)
        
    mean_vort = torch.tensor(mean_vort)[:,None]
    var_vort = torch.tensor(var_vort)[:,None]
    return mean_vort, var_vort
//  [python]
def __init__(self, mu, cov, device='cuda'):
        super(GaussianDistribution, self).__init__()

        self.mu = mu
        self.cov = cov
        self.precision = torch.inverse(cov)

        self.R = torch.cholesky(self.cov)
        self.normal = torch.distributions.normal.Normal(
            torch.zeros_like(mu), torch.ones_like(mu))
//  [python]
def __init__(self, mu, cov, device='cuda'):
        super(GaussianDistribution, self).__init__()

        self.mu = mu
        self.cov = cov
        self.precision = torch.inverse(cov)

        self.R = torch.linalg.cholesky(self.cov)
        self.normal = torch.distributions.normal.Normal(
            torch.zeros_like(mu), torch.ones_like(mu))
//  [python]
def __init__(
        self,
        sample_config,
    ):
        super().__init__(**vars(sample_config))

        self.Z = torch.zeros(
            self.horizon * self.d_action,
            dtype=self.tensor_args.dtype,
            device=self.tensor_args.device,
        )
        if self.scale_tril is None and self.covariance_matrix is not None:
            self.scale_tril = self.tensor_args.to_device(
                torch.cholesky(covariance_matrix.to("cpu"))
            )
        self.samples = None
        self.sample_shape = 0
        self.ndims = self.horizon * self.d_action
        self.stomp_matrix, self.stomp_scale_tril = None, None
//  [python]
def __init__(
        self,
        sample_config,
    ):
        super().__init__(**vars(sample_config))

        self.Z = torch.zeros(
            self.horizon * self.d_action,
            dtype=self.tensor_args.dtype,
            device=self.tensor_args.device,
        )
        if self.scale_tril is None and self.covariance_matrix is not None:
            self.scale_tril = self.tensor_args.to_device(
                torch.linalg.cholesky(covariance_matrix.to("cpu"))
            )
        self.samples = None
        self.sample_shape = 0
        self.ndims = self.horizon * self.d_action
        self.stomp_matrix, self.stomp_scale_tril = None, None
//  [python]
def forward(self, x, x_var=0.):
        """
        Forward pass of the CIGP module.

        Args:
            x (torch.Tensor or list): Input tensor or list of tensors.
            x_var (float): Variance of the input tensor.

        Returns:
            tuple: Tuple containing the predicted mean and variance.
        """
        x = self.check_single_tensor(x)

        if self.train_x is None:
            MFGP_LOG.e("gp model model hasn't been trained. predict failed")
            return None
    
        with torch.no_grad():
            Sigma = self.kernel(self.train_x, self.train_x) + JITTER * torch.eye(self.train_x.size(0), device=list(self.parameters())[0].device)
            _noise = self.noise_box.get()
            Sigma = Sigma + _noise.pow(-1) * torch.eye(self.train_x.size(0), device=list(self.parameters())[0].device)

            L = torch.cholesky(Sigma)
            kx = self.kernel(self.train_x, x)
            LinvKx, _ = torch.triangular_solve(kx, L, upper = False)
            # LinvKx = torch.linalg.solve_triangular(kx, L, upper = False)      # torch recommand but not support nonsquare matrix

            u = kx.t() @ torch.cholesky_solve(self.train_y, L)

            var_diag = self.kernel(x, x).diag().view(-1, 1) - (LinvKx**2).sum(dim = 0).view(-1, 1)
            _noise = self.noise_box.get()
            var_diag = var_diag + _noise.pow(-1)
            var_diag = var_diag.expand_as(u)

            var_diag = var_diag + x_var

        return u, var_diag