jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def log_marginal_likelihood(self):
        K = self.get_covariance_matrix(self.x, self.x, self.dxdr, self.dxdr)
        self.L = torch.cholesky(K + self.noise * torch.eye(len(K)), upper=False)
        self.alpha = torch.cholesky_solve(self.y, self.L)

        # log marginal likelihood
        MLL = -0.5 * torch.einsum("ik,ik", self.y, self.alpha)
        MLL -= torch.log(torch.diag(self.L)).sum()
        MLL -= self.L.shape[0] / 2 * math.log(2 * math.pi)
        
        mu = K.T @ self.alpha
        error = torch.abs(mu - self.y)
        E_mae = error[:self.no_of_structures].mean()
        F_mae = error[self.no_of_structures:].mean()

        return -MLL, E_mae, F_mae
//  [python]
def log_marginal_likelihood(self):
        K = self.get_covariance_matrix(self.x, self.x, self.dxdr, self.dxdr)
        self.L = torch.linalg.cholesky(K + self.noise * torch.eye(len(K)))
        self.alpha = torch.linalg.solve_triangular(self.L, self.y, upper=False)

        # log marginal likelihood
        MLL = -0.5 * torch.einsum("ik,ik", self.y, self.alpha)
        MLL -= torch.log(torch.diag(self.L)).sum()
        MLL -= self.L.shape[0] / 2 * math.log(2 * math.pi)
        
        mu = K.T @ self.alpha
        error = torch.abs(mu - self.y)
        E_mae = error[:self.no_of_structures].mean()
        F_mae = error[self.no_of_structures:].mean()

        return -MLL, E_mae, F_mae
//  [python]
def logdet(self):
        bs,c,h,w = self._shape
        padded_weight = F.pad(self.conv.weight,(0,h-3,0,w-3))
        w_fft = torch.rfft(padded_weight, 2, onesided=False, normalized=False)
        # Lift to real valued space
        D = phi(w_fft).permute(2,3,0,1)
        Dt = D.permute(0, 1, 3, 2) #transpose of D
        lhs = torch.matmul(D, Dt)
        scale = lhs.data.norm().detach()/np.sqrt(np.prod(Dt.shape))
        chol_output = torch.cholesky(lhs+3e-5*scale*torch.eye(lhs.size(-1)).to(lhs.device))
        eigs = torch.diagonal(chol_output,dim1=-2,dim2=-1)
        logdet = (eigs.log().sum() / 2.0).expand(bs)
        # 1/4 \sum_{h,w} log det (DDt)
        return logdet
//  [python]
def logdet(self):
        bs,c,h,w = self._shape
        padded_weight = F.pad(self.conv.weight,(0,h-3,0,w-3))
        w_fft = torch.rfft(padded_weight, 2, onesided=False, normalized=False)
        # Lift to real valued space
        D = phi(w_fft).permute(2,3,0,1)
        Dt = D.permute(0, 1, 3, 2) #transpose of D
        lhs = torch.matmul(D, Dt)
        scale = lhs.data.norm().detach()/np.sqrt(np.prod(Dt.shape))
        chol_output = torch.linalg.cholesky(lhs+3e-5*scale*torch.eye(lhs.size(-1)).to(lhs.device))
        eigs = torch.diagonal(chol_output,dim1=-2,dim2=-1)
        logdet = (eigs.log().sum() / 2.0).expand(bs)
        # 1/4 \sum_{h,w} log det (DDt)
        return logdet
//  [python]
def reduce_func_singular_values(self,func):
        bs,c,h,w = self._shape
        padded_weight = F.pad(self.conv.weight,(0,h-3,0,w-3))
        w_fft = torch.rfft(padded_weight, 2, onesided=False, normalized=False)
        # Lift to real valued space
        D = phi(w_fft).permute(2,3,0,1)
        Dt = D.permute(0, 1, 3, 2) #transpose of D
        lhs = torch.matmul(D, Dt)
        scale = lhs.data.norm().detach()/np.sqrt(np.prod(Dt.shape))
        chol_output = torch.cholesky(lhs+1e-4*scale*torch.eye(lhs.size(-1)).to(lhs.device))
        eigs = torch.diagonal(chol_output,dim1=-2,dim2=-1)
        logdet = (func(eigs).sum() / 2.0).expand(bs)
        # 1/4 \sum_{h,w} log det (DDt)
        return logdet
//  [python]
def reduce_func_singular_values(self,func):
        bs,c,h,w = self._shape
        padded_weight = F.pad(self.conv.weight,(0,h-3,0,w-3))
        w_fft = torch.rfft(padded_weight, 2, onesided=False, normalized=False)
        # Lift to real valued space
        D = phi(w_fft).permute(2,3,0,1)
        Dt = D.permute(0, 1, 3, 2) #transpose of D
        lhs = torch.matmul(D, Dt)
        scale = lhs.data.norm().detach()/np.sqrt(np.prod(Dt.shape))
        chol_output = torch.linalg.cholesky(lhs+1e-4*scale*torch.eye(lhs.size(-1)).to(lhs.device))
        eigs = torch.diagonal(chol_output,dim1=-2,dim2=-1)
        logdet = (func(eigs).sum() / 2.0).expand(bs)
        # 1/4 \sum_{h,w} log det (DDt)
        return logdet
//  [python]
def estimate_variance_batch(self, X, criterion, dist, var0):
        tau = 1/var0
        opt = KFAC(self.net)
        U_halfs = {}
        V_halfs = {}

        # Populate parameters with the means
        self.sample(scale=0, require_grad=True)

        output = self(X)
        distribution = dist(logits=output)
        y = distribution.sample()
        loss = criterion(output, y)
        loss.backward()
        opt.step()

        with torch.no_grad():
            for group in opt.param_groups:
                if len(group['params']) == 2:
                    weight, bias = group['params']
                else:
                    weight = group['params'][0]
                    bias = None

                module = group['mod']
                state = opt.state[module]

                U = state['ggt']
                V = state['xxt']

                m, n = int(U.shape[0]), int(V.shape[0])

                # Add priors
                n_data = len(X)
                U = sqrt(n_data)*U + sqrt(tau)*torch.eye(m)
                V = sqrt(n_data)*V + sqrt(tau)*torch.eye(n)

                U_half = torch.cholesky(torch.inverse(U), upper=False)
                V_half = torch.cholesky(torch.inverse(V), upper=True)

                U_halfs[(module, 'weight')] = U_half
                V_halfs[(module, 'weight')] = V_half


        for module, name in self.params:
            mod_class = module.__class__.__name__
            if mod_class not in ['Linear', 'Conv2d']:
                continue

            if name == 'bias':
                continue

            U_half = U_halfs[(module, name)]
            V_half = V_halfs[(module, name)]

            module.__getattr__(f'{name}_U_half').copy_(U_half)
            module.__getattr__(f'{name}_V_half').copy_(V_half)
//  [python]
def estimate_variance_batch(self, X, criterion, dist, var0):
        tau = 1/var0
        opt = KFAC(self.net)
        U_halfs = {}
        V_halfs = {}

        # Populate parameters with the means
        self.sample(scale=0, require_grad=True)

        output = self(X)
        distribution = dist(logits=output)
        y = distribution.sample()
        loss = criterion(output, y)
        loss.backward()
        opt.step()

        with torch.no_grad():
            for group in opt.param_groups:
                if len(group['params']) == 2:
                    weight, bias = group['params']
                else:
                    weight = group['params'][0]
                    bias = None

                module = group['mod']
                state = opt.state[module]

                U = state['ggt']
                V = state['xxt']

                m, n = int(U.shape[0]), int(V.shape[0])

                # Add priors
                n_data = len(X)
                U = sqrt(n_data)*U + sqrt(tau)*torch.eye(m)
                V = sqrt(n_data)*V + sqrt(tau)*torch.eye(n)

                U_half = torch.linalg.cholesky(torch.inverse(U), upper=False)
                V_half = torch.linalg.cholesky(torch.inverse(V), upper=True)

                U_halfs[(module, 'weight')] = U_half
                V_halfs[(module, 'weight')] = V_half


        for module, name in self.params:
            mod_class = module.__class__.__name__
            if mod_class not in ['Linear', 'Conv2d']:
                continue

            if name == 'bias':
                continue

            U_half = U_halfs[(module, name)]
            V_half = V_halfs[(module, name)]

            module.__getattr__(f'{name}_U_half').copy_(U_half)
            module.__getattr__(f'{name}_V_half').copy_(V_half)
//  [python]
def estimate_variance(self, train_loader, var0):
        tau = 1/var0
        criterion = nn.CrossEntropyLoss()
        opt = KFAC(self.net)
        U_halfs = {}
        V_halfs = {}

        # Populate parameters with the means
        self.sample(scale=0, require_grad=True)

        # for x, y in tqdm(train_loader):
        for x, y in train_loader:
            x = x.cuda(non_blocking=True)

            output = self(x)
            distribution = torch.distributions.Categorical(logits=output)
            y = distribution.sample()

            loss = criterion(output, y)
            loss.backward()
            opt.step()

        with torch.no_grad():
            for group in opt.param_groups:
                if len(group['params']) == 2:
                    weight, bias = group['params']
                else:
                    weight = group['params'][0]
                    bias = None

                module = group['mod']
                state = opt.state[module]

                U = state['ggt']
                V = state['xxt']

                m, n = int(U.shape[0]), int(V.shape[0])

                # Add priors
                n_data = len(train_loader.dataset)
                U = sqrt(n_data)*U + sqrt(tau)*torch.eye(m, device='cuda')
                V = sqrt(n_data)*V + sqrt(tau)*torch.eye(n, device='cuda')

                U_half = torch.cholesky(torch.inverse(U), upper=False)
                V_half = torch.cholesky(torch.inverse(V), upper=True)

                U_halfs[(module, 'weight')] = U_half
                V_halfs[(module, 'weight')] = V_half


        for module, name in self.params:
            mod_class = module.__class__.__name__
            if mod_class not in ['Linear', 'Conv2d']:
                continue

            if name == 'bias':
                continue

            U_half = U_halfs[(module, name)]
            V_half = V_halfs[(module, name)]

            module.__getattr__(f'{name}_U_half').copy_(U_half)
            module.__getattr__(f'{name}_V_half').copy_(V_half)
//  [python]
def estimate_variance(self, train_loader, var0):
        tau = 1/var0
        criterion = nn.CrossEntropyLoss()
        opt = KFAC(self.net)
        U_halfs = {}
        V_halfs = {}

        # Populate parameters with the means
        self.sample(scale=0, require_grad=True)

        # for x, y in tqdm(train_loader):
        for x, y in train_loader:
            x = x.cuda(non_blocking=True)

            output = self(x)
            distribution = torch.distributions.Categorical(logits=output)
            y = distribution.sample()

            loss = criterion(output, y)
            loss.backward()
            opt.step()

        with torch.no_grad():
            for group in opt.param_groups:
                if len(group['params']) == 2:
                    weight, bias = group['params']
                else:
                    weight = group['params'][0]
                    bias = None

                module = group['mod']
                state = opt.state[module]

                U = state['ggt']
                V = state['xxt']

                m, n = int(U.shape[0]), int(V.shape[0])

                # Add priors
                n_data = len(train_loader.dataset)
                U = sqrt(n_data)*U + sqrt(tau)*torch.eye(m, device='cuda')
                V = sqrt(n_data)*V + sqrt(tau)*torch.eye(n, device='cuda')

                U_half = torch.linalg.cholesky(torch.inverse(U), upper=False)
                V_half = torch.linalg.cholesky(torch.inverse(V), upper=True)

                U_halfs[(module, 'weight')] = U_half
                V_halfs[(module, 'weight')] = V_half


        for module, name in self.params:
            mod_class = module.__class__.__name__
            if mod_class not in ['Linear', 'Conv2d']:
                continue

            if name == 'bias':
                continue

            U_half = U_halfs[(module, name)]
            V_half = V_halfs[(module, name)]

            module.__getattr__(f'{name}_U_half').copy_(U_half)
            module.__getattr__(f'{name}_V_half').copy_(V_half)
//  [python]
def test_cholesky(N0=23):
    tmp0 = np.random.randn(N0,N0)
    np0 = np.matmul(tmp0, tmp0.T) + np.eye(N0)*0.1
    ret_ = np.linalg.cholesky(np0)
    ret0 = torch.cholesky(torch.tensor(np0)).numpy()
    assert hfe(ret_, ret0) < 1e-7
    if torch.cuda.is_available():
        tmp0 = torch.tensor(np0, device='cuda')
        ret1 = torch.cholesky(tmp0).to('cpu').numpy()
        assert hfe(ret_, ret1) < 1e-7
//  [python]
def test_cholesky(N0=23):
    tmp0 = np.random.randn(N0,N0)
    np0 = np.matmul(tmp0, tmp0.T) + np.eye(N0)*0.1
    ret_ = np.linalg.cholesky(np0)
    ret0 = torch.linalg.cholesky(torch.tensor(np0)).numpy()
    assert hfe(ret_, ret0) < 1e-7
    if torch.cuda.is_available():
        tmp0 = torch.tensor(np0, device='cuda')
        ret1 = torch.linalg.cholesky(tmp0).to('cpu').numpy()
        assert hfe(ret_, ret1) < 1e-7
//  [python]
def forward(ctx, Q, P):
        L = torch.cholesky(Q)
        D = torch.cholesky_solve(P, L)  # D = Q-1 @ P
        ctx.save_for_backward(L, D)
        return D
//  [python]
def forward(ctx, Q, P):
        L = torch.linalg.cholesky(Q)
        D = torch.cholesky_solve(P, L)  # D = Q-1 @ P
        ctx.save_for_backward(L, D)
        return D
//  [python]
def test_cholesky_solve(N0=23, N1=13):
    tmp0 = np.random.randn(N0,N0)
    np0 = np.matmul(tmp0, tmp0.T) + np.eye(N0)*0.1
    np1 = np.random.randn(N0,N1)
    ret_ = np.linalg.solve(np0, np1)
    tmp0 = torch.cholesky(torch.tensor(np0))
    ret0 = torch.cholesky_solve(torch.tensor(np1), tmp0).numpy()
    assert hfe(ret_, ret0) < 1e-7
    if torch.cuda.is_available():
        tmp0 = torch.tensor(np0, device='cuda')
        tmp1 = torch.tensor(np1, device='cuda')
        ret1 = torch.cholesky_solve(tmp1, torch.cholesky(tmp0)).to('cpu').numpy()
        assert hfe(ret_, ret1) < 1e-7
//  [python]
def test_cholesky_solve(N0=23, N1=13):
    tmp0 = np.random.randn(N0,N0)
    np0 = np.matmul(tmp0, tmp0.T) + np.eye(N0)*0.1
    np1 = np.random.randn(N0,N1)
    ret_ = np.linalg.solve(np0, np1)
    tmp0 = torch.linalg.cholesky(torch.tensor(np0))
    ret0 = torch.cholesky_solve(torch.tensor(np1), tmp0).numpy()
    assert hfe(ret_, ret0) < 1e-7
    if torch.cuda.is_available():
        tmp0 = torch.tensor(np0, device='cuda')
        tmp1 = torch.tensor(np1, device='cuda')
        ret1 = torch.cholesky_solve(tmp1, torch.linalg.cholesky(tmp0)).to('cpu').numpy()
        assert hfe(ret_, ret1) < 1e-7
//  [python]
def demo_compare_torch_cupy_SPD_inverse_speed():
    N0 = 512
    N1 = 100
    tmp0 = np.random.randn(N0,N0)
    np0 = (np.matmul(tmp0, tmp0.T)/100 + np.eye(N0)*0.1).astype(np.float32)
    torch0 = torch.tensor(np0, device=torch.device('cuda:0'))
    t0 = time.time()
    for _ in range(N1):
        _ = torch.cholesky(torch0)
    t_torch = (time.time()-t0) / N1

    t0 = time.time()
    for _ in range(N1):
        _ = torch_SPD_inverse(torch0)
    t_cupy = (time.time()-t0) / N1
    print(f'[N0={N0}] time(torch)={t_torch}')
    print(f'[N0={N0}] time(cupy) ={t_cupy}')
//  [python]
def demo_compare_torch_cupy_SPD_inverse_speed():
    N0 = 512
    N1 = 100
    tmp0 = np.random.randn(N0,N0)
    np0 = (np.matmul(tmp0, tmp0.T)/100 + np.eye(N0)*0.1).astype(np.float32)
    torch0 = torch.tensor(np0, device=torch.device('cuda:0'))
    t0 = time.time()
    for _ in range(N1):
        _ = torch.linalg.cholesky(torch0)
    t_torch = (time.time()-t0) / N1

    t0 = time.time()
    for _ in range(N1):
        _ = torch_SPD_inverse(torch0)
    t_cupy = (time.time()-t0) / N1
    print(f'[N0={N0}] time(torch)={t_torch}')
    print(f'[N0={N0}] time(cupy) ={t_cupy}')
//  [python]
def _initalize_parameters(self, pi: torch.Tensor, loc: torch.Tensor, cov: torch.Tensor):
        # resize and copy inplace
        self.pi = torch.nn.Parameter(pi)
        self.loc = torch.nn.Parameter(loc)

        # keep diagonal in log space
        if self.covariance_type == "diag":
            # obtain cholesky decomposition of the diagonal
            # convert to diag matrix, obtain decomposition, only keep diagonal
            before = cov.shape
            cov = torch.stack([torch.diagonal(torch.cholesky(torch.diag(cov[i])))
                               for i in range(self.k)])

            assert cov.shape == before

            # keep in log space
            cov = torch.log(cov)
        self.cov = torch.nn.Parameter(cov)

        self._fitted = True
//  [python]
def _initalize_parameters(self, pi: torch.Tensor, loc: torch.Tensor, cov: torch.Tensor):
        # resize and copy inplace
        self.pi = torch.nn.Parameter(pi)
        self.loc = torch.nn.Parameter(loc)

        # keep diagonal in log space
        if self.covariance_type == "diag":
            # obtain cholesky decomposition of the diagonal
            # convert to diag matrix, obtain decomposition, only keep diagonal
            before = cov.shape
            cov = torch.stack([torch.diagonal(torch.linalg.cholesky(torch.diag(cov[i])))
                               for i in range(self.k)])

            assert cov.shape == before

            # keep in log space
            cov = torch.log(cov)
        self.cov = torch.nn.Parameter(cov)

        self._fitted = True
//  [python]
def forward(self, Xte):
        n_test = Xte.size(0)
        Xte = ( Xte - self.Xmean.expand_as(Xte) ) / self.Xstd.expand_as(Xte)

        Sigma = self.kernel(self.X, self.X) + self.log_beta.exp().pow(-1) * torch.eye(self.X.size(0)) \
            + JITTER * torch.eye(self.X.size(0))

        kx = self.kernel(self.X, Xte)
        L = torch.cholesky(Sigma)
        LinvKx,_ = torch.triangular_solve(kx, L, upper = False)

        # option 1
        mean = self.meanFunc(Xte)
        # mean = mean + kx.t() @ torch.cholesky_solve(self.Y, L)  # torch.linalg.cholesky()
        mean = mean + kx.t() @ torch.cholesky_solve(self.Y-self.meanFunc(self.X), L)  # torch.linalg.cholesky()

        
        var_diag = self.kernel(Xte, Xte).diag().view(-1, 1) \
            - (LinvKx**2).sum(dim = 0).view(-1, 1)

        # add the noise uncertainty
        var_diag = var_diag + self.log_beta.exp().pow(-1)

        # de-normalized
        mean = mean * self.Ystd.expand_as(mean) + self.Ymean.expand_as(mean)
        var_diag = var_diag.expand_as(mean) * self.Ystd**2

        return mean, var_diag
//  [python]
def forward(self, Xte):
        n_test = Xte.size(0)
        Xte = ( Xte - self.Xmean.expand_as(Xte) ) / self.Xstd.expand_as(Xte)

        Sigma = self.kernel(self.X, self.X) + self.log_beta.exp().pow(-1) * torch.eye(self.X.size(0)) \
            + JITTER * torch.eye(self.X.size(0))

        kx = self.kernel(self.X, Xte)
        L = torch.linalg.cholesky(Sigma)
        LinvKx,_ = torch.triangular_solve(kx, L, upper = False)

        # option 1
        mean = self.meanFunc(Xte)
        # mean = mean + kx.t() @ torch.cholesky_solve(self.Y, L)  # torch.linalg.cholesky()
        mean = mean + kx.t() @ torch.linalg.solve_triangular(L, torch.linalg.solve_triangular(L.t(),self.Y-self.meanFunc(self.X), upper=False), upper=False)  # torch.linalg.cholesky()

        
        var_diag = self.kernel(Xte, Xte).diag().view(-1, 1) \
            - (LinvKx**2).sum(dim = 0).view(-1, 1)

        # add the noise uncertainty
        var_diag = var_diag + self.log_beta.exp().pow(-1)

        # de-normalized
        mean = mean * self.Ystd.expand_as(mean) + self.Ymean.expand_as(mean)
        var_diag = var_diag.expand_as(mean) * self.Ystd**2

        return mean, var_diag
//  [python]
def get_R0(self, forecaster, Y_train, covariates_full, sample_size, batch_size):
        n_batch = sample_size // batch_size
        R_list = []
        map_list = []

        for i in range(n_batch):
            with torch.no_grad():
                with pyro.plate("particles", batch_size, dim=-3):
                    map_estimates = forecaster.guide(Y_train[:1, :], covariates_full)
                map_list.append(map_estimates)
                data_dim = Y_train.size(-1)
                time_dim = covariates_full.size(-2)

                kernel_lengthscale = map_estimates['r0_lengthscale']
                kernel_var = map_estimates['r0_kernel_var']
                covariates_pyro = pyro_model.helper.reshape_covariates_pyro(covariates_full,
                                                                            self.n_country).to(torch.float)

                kernel_lengthscale = kernel_lengthscale.transpose(-1, -2)
                kernel_var = kernel_var.transpose(-1, -2)

                # covariate
                covariates_pyro = covariates_pyro.transpose(1, 0)
                d_times_t = data_dim * time_dim
                # dxt, 1, p
                covariates_pyro = covariates_pyro.reshape(d_times_t, 1, covariates_pyro.size(-1))

                var = pyro_model.helper.tensor_RBF(kernel_lengthscale, kernel_var, covariates_pyro)
                var = var + torch.eye(var.shape[-1]) * 0.01
                A = torch.cholesky(var)

                iid_n = map_estimates['r0_iid_n']
                iid_n = torch.cat([iid_n, torch.randn(iid_n.size(0), iid_n.size(1), time_dim - iid_n.size(-1))], dim=-1)
                iid_n = iid_n.unsqueeze(-1)
                iid_n = iid_n.reshape(iid_n.size(0), 1, d_times_t, 1)

                weight = torch.sigmoid(torch.einsum('abij,abjk->abik', A, iid_n))
                weight = weight[:, 0, ...]  # get rid of batch dimension
                weight = weight.reshape(weight.size(0), data_dim, time_dim, 1)
                R00 = map_estimates['R00']
                R0 = R00 * weight[..., 0]
                R_list.append(R0)

        R0 = torch.cat(R_list, dim=0)
        R0low, R0mid, R0high = quantile(R0, (0.1, 0.5, 0.9), dim=0).squeeze(-1)

        map_estimates = dict.fromkeys(map_list[0].keys())
        for k in map_list[0].keys():
            k_list = []
            for m in map_list:
                k_list.append(m[k])
            kest = torch.cat(k_list, dim=0)
            map_estimates[k] = kest

        return R0low, R0mid, R0high, map_estimates
//  [python]
def get_R0(self, forecaster, Y_train, covariates_full, sample_size, batch_size):
        n_batch = sample_size // batch_size
        R_list = []
        map_list = []

        for i in range(n_batch):
            with torch.no_grad():
                with pyro.plate("particles", batch_size, dim=-3):
                    map_estimates = forecaster.guide(Y_train[:1, :], covariates_full)
                map_list.append(map_estimates)
                data_dim = Y_train.size(-1)
                time_dim = covariates_full.size(-2)

                kernel_lengthscale = map_estimates['r0_lengthscale']
                kernel_var = map_estimates['r0_kernel_var']
                covariates_pyro = pyro_model.helper.reshape_covariates_pyro(covariates_full,
                                                                            self.n_country).to(torch.float)

                kernel_lengthscale = kernel_lengthscale.transpose(-1, -2)
                kernel_var = kernel_var.transpose(-1, -2)

                # covariate
                covariates_pyro = covariates_pyro.transpose(1, 0)
                d_times_t = data_dim * time_dim
                # dxt, 1, p
                covariates_pyro = covariates_pyro.reshape(d_times_t, 1, covariates_pyro.size(-1))

                var = pyro_model.helper.tensor_RBF(kernel_lengthscale, kernel_var, covariates_pyro)
                var = var + torch.eye(var.shape[-1]) * 0.01
                A = torch.linalg.cholesky(var)

                iid_n = map_estimates['r0_iid_n']
                iid_n = torch.cat([iid_n, torch.randn(iid_n.size(0), iid_n.size(1), time_dim - iid_n.size(-1))], dim=-1)
                iid_n = iid_n.unsqueeze(-1)
                iid_n = iid_n.reshape(iid_n.size(0), 1, d_times_t, 1)

                weight = torch.sigmoid(torch.einsum('abij,abjk->abik', A, iid_n))
                weight = weight[:, 0, ...]  # get rid of batch dimension
                weight = weight.reshape(weight.size(0), data_dim, time_dim, 1)
                R00 = map_estimates['R00']
                R0 = R00 * weight[..., 0]
                R_list.append(R0)

        R0 = torch.cat(R_list, dim=0)
        R0low, R0mid, R0high = quantile(R0, (0.1, 0.5, 0.9), dim=0).squeeze(-1)

        map_estimates = dict.fromkeys(map_list[0].keys())
        for k in map_list[0].keys():
            k_list = []
            for m in map_list:
                k_list.append(m[k])
            kest = torch.cat(k_list, dim=0)
            map_estimates[k] = kest

        return R0low, R0mid, R0high, map_estimates
//  [python]
def forward(self, x_tr, y_tr, index, Xte):
        n_test = Xte.size(0)
        Xte = ( Xte - self.Xmean.expand_as(Xte) ) / self.Xstd.expand_as(Xte)

        Sigma = self.kernel(x_tr, x_tr) + self.log_beta.exp().pow(-1) * torch.eye(x_tr.size(0)) \
            + JITTER * torch.eye(x_tr.size(0))

        kx = self.kernel(x_tr, Xte)
        L = torch.cholesky(Sigma)
        LinvKx,_ = torch.triangular_solve(kx, L, upper = False)

        # option 1
        mean = kx.t() @ torch.cholesky_solve(y_tr, L)  # torch.linalg.cholesky()
        
        var_diag = self.kernel(Xte, Xte).diag().view(-1, 1) \
            - (LinvKx**2).sum(dim = 0).view(-1, 1)

        # add the noise uncertainty
        var_diag = var_diag + self.log_beta.exp().pow(-1)

        # de-normalized
        # mean = mean * self.Ystd[index].expand_as(mean) + self.Ymean[index].expand_as(mean)
        # var_diag = var_diag.expand_as(mean) * self.Ystd[index]**2

        return mean, var_diag
//  [python]
def forward(self, x_tr, y_tr, index, Xte):
        n_test = Xte.size(0)
        Xte = ( Xte - self.Xmean.expand_as(Xte) ) / self.Xstd.expand_as(Xte)

        Sigma = self.kernel(x_tr, x_tr) + self.log_beta.exp().pow(-1) * torch.eye(x_tr.size(0)) \
            + JITTER * torch.eye(x_tr.size(0))

        kx = self.kernel(x_tr, Xte)
        L = torch.linalg.cholesky(Sigma)
        LinvKx,_ = torch.triangular_solve(kx, L, upper = False)

        # option 1
        mean = kx.t() @ torch.linalg.cholesky_solve(y_tr, L)  # torch.linalg.cholesky()
        
        var_diag = self.kernel(Xte, Xte).diag().view(-1, 1) \
            - (LinvKx**2).sum(dim = 0).view(-1, 1)

        # add the noise uncertainty
        var_diag = var_diag + self.log_beta.exp().pow(-1)

        # de-normalized
        # mean = mean * self.Ystd[index].expand_as(mean) + self.Ymean[index].expand_as(mean)
        # var_diag = var_diag.expand_as(mean) * self.Ystd[index]**2

        return mean, var_diag
//  [python]
def forward_res(self, x_tr, y_tr, l1, h1, l2, h2, index, Xte):
        n_test = Xte.size(0)
        Xte = ( Xte - self.Xmean.expand_as(Xte) ) / self.Xstd.expand_as(Xte)

        Sigma = self.kernel_res(x_tr, x_tr, l1, h1, l2, h2) + self.log_beta.exp().pow(-1) * torch.eye(x_tr.size(0)) \
            + JITTER * torch.eye(x_tr.size(0))

        kx = self.kernel_res(x_tr, Xte, l1, h1, l2, h2)
        L = torch.cholesky(Sigma)
        LinvKx,_ = torch.triangular_solve(kx, L, upper = False)

        # option 1
        y_res = y_tr[0] - (-self.b * (h1 - l1)).exp() * y_tr[1]
        mean = kx.t() @ torch.cholesky_solve(y_res, L)  # torch.linalg.cholesky()
        
        var_diag = self.kernel_res(Xte, Xte, l1, h1, l2, h2).diag().view(-1, 1) \
            - (LinvKx**2).sum(dim = 0).view(-1, 1)

        # add the noise uncertainty
        var_diag = var_diag + self.log_beta.exp().pow(-1)

        # de-normalized
        # mean = mean * self.Ystd[index].expand_as(mean) + self.Ymean[index].expand_as(mean)
        # var_diag = var_diag.expand_as(mean) * self.Ystd[index]**2
        
        return mean, var_diag
//  [python]
def forward_res(self, x_tr, y_tr, l1, h1, l2, h2, index, Xte):
        n_test = Xte.size(0)
        Xte = ( Xte - self.Xmean.expand_as(Xte) ) / self.Xstd.expand_as(Xte)

        Sigma = self.kernel_res(x_tr, x_tr, l1, h1, l2, h2) + self.log_beta.exp().pow(-1) * torch.eye(x_tr.size(0)) \
            + JITTER * torch.eye(x_tr.size(0))

        kx = self.kernel_res(x_tr, Xte, l1, h1, l2, h2)
        L = torch.linalg.cholesky(Sigma)
        LinvKx,_ = torch.triangular_solve(kx, L, upper = False)

        # option 1
        y_res = y_tr[0] - (-self.b * (h1 - l1)).exp() * y_tr[1]
        mean = kx.t() @ torch.linalg.solve(L.t(), torch.linalg.solve(L, y_res))  # torch.linalg.cholesky()
        
        var_diag = self.kernel_res(Xte, Xte, l1, h1, l2, h2).diag().view(-1, 1) \
            - (LinvKx**2).sum(dim = 0).view(-1, 1)

        # add the noise uncertainty
        var_diag = var_diag + self.log_beta.exp().pow(-1)

        # de-normalized
        # mean = mean * self.Ystd[index].expand_as(mean) + self.Ymean[index].expand_as(mean)
        # var_diag = var_diag.expand_as(mean) * self.Ystd[index]**2
        
        return mean, var_diag
//  [python]
def _get_kernel_and_auxiliary_variables(
        self, x_tr, y_tr, x_te=None,
    ):
        """ Get kernel and auxiliary variables for forward pass. """

        # compute the kernels
        k_tr_tr = self._perturb(self.kernel.forward(x_tr, x_tr))

        if x_te is not None:  # during test
            k_te_te = self._perturb(self.kernel.forward(x_te, x_te))
            k_te_tr = self._perturb(self.kernel.forward(x_te, x_tr))
            # k_tr_te = self.forward(x_tr, x_te)
            k_tr_te = k_te_tr.t()  # save time

        else:  # during train
            k_te_te = k_te_tr = k_tr_te = k_tr_tr

        # (batch_size_tr, batch_size_tr)
        k_plus_sigma = k_tr_tr + torch.exp(self.log_sigma) * torch.eye(
            k_tr_tr.shape[0], device=k_tr_tr.device
        )

        # (batch_size_tr, batch_size_tr)
        l_low = torch.cholesky(k_plus_sigma)
        l_up = l_low.t()

        # (batch_size_tr. 1)
        l_low_over_y, _ = torch.triangular_solve(
            input=y_tr, A=l_low, upper=False
        )

        # (batch_size_tr, 1)
        alpha, _ = torch.triangular_solve(
            input=l_low_over_y, A=l_up, upper=True
        )

        return k_tr_tr, k_te_te, k_te_tr, k_tr_te, l_low, alpha
//  [python]
def _get_kernel_and_auxiliary_variables(
        self, x_tr, y_tr, x_te=None,
    ):
        """ Get kernel and auxiliary variables for forward pass. """

        # compute the kernels
        k_tr_tr = self._perturb(self.kernel.forward(x_tr, x_tr))

        if x_te is not None:  # during test
            k_te_te = self._perturb(self.kernel.forward(x_te, x_te))
            k_te_tr = self._perturb(self.kernel.forward(x_te, x_tr))
            # k_tr_te = self.forward(x_tr, x_te)
            k_tr_te = k_te_tr.t()  # save time

        else:  # during train
            k_te_te = k_te_tr = k_tr_te = k_tr_tr

        # (batch_size_tr, batch_size_tr)
        k_plus_sigma = k_tr_tr + torch.exp(self.log_sigma) * torch.eye(
            k_tr_tr.shape[0], device=k_tr_tr.device
        )

        # (batch_size_tr, batch_size_tr)
        l_low = torch.linalg.cholesky(k_plus_sigma)
        l_up = l_low.t()

        # (batch_size_tr. 1)
        l_low_over_y, _ = torch.triangular_solve(
            input=y_tr, A=l_low, upper=False
        )

        # (batch_size_tr, 1)
        alpha, _ = torch.triangular_solve(
            input=l_low_over_y, A=l_up, upper=True
        )

        return k_tr_tr, k_te_te, k_te_tr, k_tr_te, l_low, alpha
//  [python]
def cholesky_inverse(input, upper=False, out=None, inplace=True):
    u = torch.cholesky(input, upper)
    u = torch.linalg.triangular_solve(u, torch.eye(u.shape[-1]), upper=upper)
    if len(u.shape) == 2:
        uit = u.T
    elif len(u.shape) == 3:
        uit = torch.permute(u, perm=(0, 2, 1))
    elif len(u.shape) == 4:
        uit = torch.permute(u, perm=(0, 1, 3, 2))
    if inplace:
        input = u @uit if upper else uit @u
        return input
    else:
        out = u @uit if upper else uit @u
        return out
//  [python]
def cholesky_inverse(input, upper=False, out=None, inplace=True):
    u = torch.linalg.cholesky(input, upper=upper)
    u = torch.linalg.triangular_solve(u, torch.eye(u.shape[-1]), upper=upper)
    if len(u.shape) == 2:
        uit = u.T
    elif len(u.shape) == 3:
        uit = torch.permute(u, perm=(0, 2, 1))
    elif len(u.shape) == 4:
        uit = torch.permute(u, perm=(0, 1, 3, 2))
    if inplace:
        input = u @uit if upper else uit @u
        return input
    else:
        out = u @uit if upper else uit @u
        return out
//  [python]
def negELBO(self):
        K_u_u = self.kernel.matrix(self.X, torch.exp(self.log_ls_h))
        K_u1 = self.kernel.cross(self.X_col, self.X, torch.exp(self.log_ls_h))
        K_u2 = self.kernel.cross(self.X_col, self.X_col, torch.exp(self.log_ls_h))
        K_u_tt_u_tt = self.kernel.matrix_pp(K_u2, self.X_col, 0, torch.exp(self.log_ls_h))
        K_u_tt = self.kernel.cross_pp_0(K_u1, self.X_col, self.X, 0, torch.exp(self.log_ls_h))
        K_u_t_u_t = self.kernel.matrix_p(K_u2, self.X_col, 0, torch.exp(self.log_ls_h))
        K_u_t = self.kernel.cross_p_0(K_u1, self.X_col, self.X, 0, torch.exp(self.log_ls_h))
        K_u_t_u_tt = self.kernel.cross_p_pp(K_u2, self.X_col, self.X_col, 0, torch.exp(self.log_ls_h))
        K = torch.cat(
            (
                torch.cat((K_u_u, K_u_tt.T, K_u_t.T), 1),
                torch.cat((K_u_tt, K_u_tt_u_tt, K_u_t_u_tt.T), 1),
                torch.cat((K_u_t, K_u_t_u_tt, K_u_t_u_t), 1),
            ),
            0,
        )
        A = torch.cholesky(K)
        self.L = torch.diag(self.L1.exp().view(-1)) + torch.matmul(self.L2, self.L2.T)
        self.L = torch.cholesky(self.L)
        Ltril = torch.tril(self.L)
        kl = self.KL_term(self.mu.clone().detach(), Ltril, torch.eye(self.M, device=self.device))
        eta = self.mu + torch.matmul(Ltril, torch.randn(self.M, 1, device=self.device))
        f = torch.matmul(A, eta)
        elbo = (-kl + 0.5 * self.N * self.log_tau - 0.5 * torch.exp(self.log_tau) * torch.sum(torch.square(f[0:self.N].view(-1, 1) - self.y)) + 0.5 * self.N_col * self.log_v - 0.5 * torch.exp(self.log_v) *
                torch.sum(torch.square(f[self.N_con:self.N_con + self.N_col].view(-1, 1) + torch.sin(f[self.N:self.N_con]).view(-1, 1) + torch.exp(self.B) * f[self.N_con + self.N_col:self.N_con + 2 * self.N_col].view(-1, 1))))
        return -elbo
//  [python]
def negELBO(self):
        K_u_u = self.kernel.matrix(self.X, torch.exp(self.log_ls_h))
        K_u1 = self.kernel.cross(self.X_col, self.X, torch.exp(self.log_ls_h))
        K_u2 = self.kernel.cross(self.X_col, self.X_col, torch.exp(self.log_ls_h))
        K_u_tt_u_tt = self.kernel.matrix_pp(K_u2, self.X_col, 0, torch.exp(self.log_ls_h))
        K_u_tt = self.kernel.cross_pp_0(K_u1, self.X_col, self.X, 0, torch.exp(self.log_ls_h))
        K_u_t_u_t = self.kernel.matrix_p(K_u2, self.X_col, 0, torch.exp(self.log_ls_h))
        K_u_t = self.kernel.cross_p_0(K_u1, self.X_col, self.X, 0, torch.exp(self.log_ls_h))
        K_u_t_u_tt = self.kernel.cross_p_pp(K_u2, self.X_col, self.X_col, 0, torch.exp(self.log_ls_h))
        K = torch.cat(
            (
                torch.cat((K_u_u, K_u_tt.T, K_u_t.T), 1),
                torch.cat((K_u_tt, K_u_tt_u_tt, K_u_t_u_tt.T), 1),
                torch.cat((K_u_t, K_u_t_u_tt, K_u_t_u_t), 1),
            ),
            0,
        )
        A = torch.linalg.cholesky(K)
        self.L = torch.diag(self.L1.exp().view(-1)) + torch.matmul(self.L2, self.L2.T)
        self.L = torch.linalg.cholesky(self.L)
        Ltril = torch.tril(self.L)
        kl = self.KL_term(self.mu.clone().detach(), Ltril, torch.eye(self.M, device=self.device))
        eta = self.mu + torch.matmul(Ltril, torch.randn(self.M, 1, device=self.device))
        f = torch.matmul(A, eta)
        elbo = (-kl + 0.5 * self.N * self.log_tau - 0.5 * torch.exp(self.log_tau) * torch.sum(torch.square(f[0:self.N].view(-1, 1) - self.y)) + 0.5 * self.N_col * self.log_v - 0.5 * torch.exp(self.log_v) *
                torch.sum(torch.square(f[self.N_con:self.N_con + self.N_col].view(-1, 1) + torch.sin(f[self.N:self.N_con]).view(-1, 1) + torch.exp(self.B) * f[self.N_con + self.N_col:self.N_con + 2 * self.N_col].view(-1, 1))))
        return -elbo
//  [python]
def forward(self, Xte, ytr_m, ytr_v, yte_m, yte_v):
        n_test = Xte.size(0)

        Xte = ( Xte - self.Xmean.expand_as(Xte) ) / (self.Xstd.expand_as(Xte) + EPS)
        ytr_m = (ytr_m - self.Ymean.expand_as(ytr_m)) / (self.Ystd.expand_as(ytr_m) + EPS)
        yte_m = (yte_m - self.Ymean.expand_as(yte_m)) / (self.Ystd.expand_as(yte_m) + EPS)
        
        Sigma = self.kernel(self.X, self.X) + self.log_beta.exp().pow(-1) * torch.eye(self.X.size(0)) \
                + JITTER * torch.eye(self.X.size(0))

        kx = self.kernel(self.X, Xte)
        L = torch.cholesky(Sigma)
        LinvKx, _ = torch.triangular_solve(kx, L, upper=False)

        # option 1
        mean = kx.t() @ torch.cholesky_solve(self.Y - ytr_m * self.rho, L)  # torch.linalg.cholesky()
        var_diag = self.log_scale.exp().expand(n_test, 1) \
                   - (LinvKx ** 2).sum(dim=0).view(-1, 1)

        # add the noise uncertainty
        var_diag = var_diag + self.log_beta.exp().pow(-1)
        # v = var_diag.shape
        
        mean = mean + yte_m * self.rho
        mean = mean * self.Ystd.expand_as(mean) + self.Ymean.expand_as(mean)

        # wtw = (self.rho.T @ self.rho).diag()
        # var_list = [ yte_v[:,0][i] * wtw + var_diag[i] * torch.ones(wtw.shape[0]) for i in range(len(yte_v[:,0]))]
        # var_diag = torch.stack(var_list)

        # var_diag = var_diag * self.Ystd ** 2

        return mean
//  [python]
def forward(self, Xte, ytr_m, ytr_v, yte_m, yte_v):
        n_test = Xte.size(0)

        Xte = ( Xte - self.Xmean.expand_as(Xte) ) / (self.Xstd.expand_as(Xte) + EPS)
        ytr_m = (ytr_m - self.Ymean.expand_as(ytr_m)) / (self.Ystd.expand_as(ytr_m) + EPS)
        yte_m = (yte_m - self.Ymean.expand_as(yte_m)) / (self.Ystd.expand_as(yte_m) + EPS)
        
        Sigma = self.kernel(self.X, self.X) + self.log_beta.exp().pow(-1) * torch.eye(self.X.size(0)) \
                + JITTER * torch.eye(self.X.size(0))

        kx = self.kernel(self.X, Xte)
        L = torch.linalg.cholesky(Sigma)
        LinvKx, _ = torch.triangular_solve(kx, L, upper=False)

        # option 1
        mean = kx.t() @ torch.linalg.solve_triangular(L, torch.linalg.solve_triangular(L.T, self.Y - ytr_m * self.rho, upper=False), upper=False)  # torch.linalg.cholesky()
        var_diag = self.log_scale.exp().expand(n_test, 1) \
                   - (LinvKx ** 2).sum(dim=0).view(-1, 1)

        # add the noise uncertainty
        var_diag = var_diag + self.log_beta.exp().pow(-1)
        # v = var_diag.shape
        
        mean = mean + yte_m * self.rho
        mean = mean * self.Ystd.expand_as(mean) + self.Ymean.expand_as(mean)

        # wtw = (self.rho.T @ self.rho).diag()
        # var_list = [ yte_v[:,0][i] * wtw + var_diag[i] * torch.ones(wtw.shape[0]) for i in range(len(yte_v[:,0]))]
        # var_diag = torch.stack(var_list)

        # var_diag = var_diag * self.Ystd ** 2

        return mean
//  [python]
def pred(self, Xte):
        K_u_u = self.kernel.matrix(self.X, torch.exp(self.log_ls_h))
        K_u1 = self.kernel.cross(self.X_col, self.X, torch.exp(self.log_ls_h))
        K_u2 = self.kernel.cross(self.X_col, self.X_col, torch.exp(self.log_ls_h))
        K_u_tt_u_tt = self.kernel.matrix_pp(K_u2, self.X_col, 0, torch.exp(self.log_ls_h))
        K_u_tt = self.kernel.cross_pp_0(K_u1, self.X_col, self.X, 0, torch.exp(self.log_ls_h))
        K_u_t_u_t = self.kernel.matrix_p(K_u2, self.X_col, 0, torch.exp(self.log_ls_h))
        K_u_t = self.kernel.cross_p_0(K_u1, self.X_col, self.X, 0, torch.exp(self.log_ls_h))
        K_u_t_u_tt = self.kernel.cross_p_pp(K_u2, self.X_col, self.X_col, 0, torch.exp(self.log_ls_h))
        K = torch.cat(
            (
                torch.cat((K_u_u, K_u_tt.T, K_u_t.T), 1),
                torch.cat((K_u_tt, K_u_tt_u_tt, K_u_t_u_tt.T), 1),
                torch.cat((K_u_t, K_u_t_u_tt, K_u_t_u_t), 1),
            ),
            0,
        )
        A = torch.cholesky(K)
        Ltril = torch.tril(self.L)
        eta = self.mu
        f2 = torch.matmul(A, eta)
        K_u1_te = self.kernel.cross(Xte, self.X, torch.exp(self.log_ls_h))
        K_u2_te = self.kernel.cross(Xte, self.X_col, torch.exp(self.log_ls_h))
        K_utt_te = self.kernel.cross_pp_0(K_u2_te.T, self.X_col, Xte, 0, torch.exp(self.log_ls_h))
        K_ut_te = self.kernel.cross_p_0(K_u2_te.T, self.X_col, Xte, 0, torch.exp(self.log_ls_h))
        K_te = torch.cat((K_u1_te, K_utt_te.T, K_ut_te.T), 1)
        pred_mean = torch.matmul(K_te, torch.linalg.solve(K, f2[0:self.N_con + 2 * self.N_col]))
        A1 = torch.cholesky(K)
        t1 = torch.linalg.solve(A1.T, Ltril[0:self.N_con + 2 * self.N_col])
        t2 = torch.matmul(t1, t1.T)
        pred_stds = (1 + self.jitter - (K_te * torch.linalg.solve(K, K_te.T).T).sum(1) + (K_te * torch.matmul(t2, K_te.T).T).sum(1))
        pred_stds = pred_stds.view(-1, 1)
        return pred_mean, pred_stds**0.5
//  [python]
def pred(self, Xte):
        K_u_u = self.kernel.matrix(self.X, torch.exp(self.log_ls_h))
        K_u1 = self.kernel.cross(self.X_col, self.X, torch.exp(self.log_ls_h))
        K_u2 = self.kernel.cross(self.X_col, self.X_col, torch.exp(self.log_ls_h))
        K_u_tt_u_tt = self.kernel.matrix_pp(K_u2, self.X_col, 0, torch.exp(self.log_ls_h))
        K_u_tt = self.kernel.cross_pp_0(K_u1, self.X_col, self.X, 0, torch.exp(self.log_ls_h))
        K_u_t_u_t = self.kernel.matrix_p(K_u2, self.X_col, 0, torch.exp(self.log_ls_h))
        K_u_t = self.kernel.cross_p_0(K_u1, self.X_col, self.X, 0, torch.exp(self.log_ls_h))
        K_u_t_u_tt = self.kernel.cross_p_pp(K_u2, self.X_col, self.X_col, 0, torch.exp(self.log_ls_h))
        K = torch.cat(
            (
                torch.cat((K_u_u, K_u_tt.T, K_u_t.T), 1),
                torch.cat((K_u_tt, K_u_tt_u_tt, K_u_t_u_tt.T), 1),
                torch.cat((K_u_t, K_u_t_u_tt, K_u_t_u_t), 1),
            ),
            0,
        )
        A = torch.linalg.cholesky(K)
        Ltril = torch.tril(self.L)
        eta = self.mu
        f2 = torch.matmul(A, eta)
        K_u1_te = self.kernel.cross(Xte, self.X, torch.exp(self.log_ls_h))
        K_u2_te = self.kernel.cross(Xte, self.X_col, torch.exp(self.log_ls_h))
        K_utt_te = self.kernel.cross_pp_0(K_u2_te.T, self.X_col, Xte, 0, torch.exp(self.log_ls_h))
        K_ut_te = self.kernel.cross_p_0(K_u2_te.T, self.X_col, Xte, 0, torch.exp(self.log_ls_h))
        K_te = torch.cat((K_u1_te, K_utt_te.T, K_ut_te.T), 1)
        pred_mean = torch.matmul(K_te, torch.linalg.solve(K, f2[0:self.N_con + 2 * self.N_col]))
        A1 = torch.linalg.cholesky(K)
        t1 = torch.linalg.solve(A1.T, Ltril[0:self.N_con + 2 * self.N_col])
        t2 = torch.matmul(t1, t1.T)
        pred_stds = (1 + self.jitter - (K_te * torch.linalg.solve(K, K_te.T).T).sum(1) + (K_te * torch.matmul(t2, K_te.T).T).sum(1))
        pred_stds = pred_stds.view(-1, 1)
        return pred_mean, pred_stds**0.5
//  [python]
def forward(self, x_te):
        """Forward pass.

        Parameters
        ----------
        x_te : `torch.Tensor`, `(x_te, hidden_dimension)`
            Test input.

        Returns
        -------
        predictive_mean : `torch.Tensor`, `(x_te, )`
            Predictive mean.

        predictive_cov : `torch.Tensor`, `(x_te, x_te)`
            Preditive covariance.

        """

        # get the kernels
        k_tr_tr, k_tr_te, k_te_tr, k_te_te = (
            self._k_tr_tr(),
            self._k_tr_te(x_te),
            self._k_te_tr(x_te),
            self._k_te_te(x_te),
        )

        # (n_tr, n_tr)
        l_k_tr_tr = torch.cholesky(k_tr_tr, upper=False)

        # (n_tr, 1)
        l_k_tr_tr_inv_mu, _ = torch.triangular_solve(
            input=self.y_tr_mu, A=l_k_tr_tr, upper=False
        )

        # (n_tr, n_tr)
        l_k_tr_tr_inv_sigma, _ = torch.triangular_solve(
            input=self._y_tr_sigma(), A=l_k_tr_tr, upper=False
        )

        # (n_tr, te)
        l_k_tr_tr_inv_k_tr_te, _ = torch.triangular_solve(
            input=k_tr_te, A=l_k_tr_tr, upper=False
        )

        # (n_te, 1)
        predictive_mean = l_k_tr_tr_inv_k_tr_te.t() @ l_k_tr_tr_inv_mu

        # (n_tr, n_te)
        k_tr_tr_inv_k_tr_te = (
            l_k_tr_tr_inv_k_tr_te.t() @ l_k_tr_tr_inv_k_tr_te
        )

        # (n_te, n_te)
        l_k_tr_tr_inv_sigma_t_at_l_k_tr_tr_inv_k_tr_te = (
            l_k_tr_tr_inv_sigma.t() @ l_k_tr_tr_inv_k_tr_te
        )

        # (n_te, n_te)
        compose_l_k_tr_tr_inv_sigma_t_at_l_k_tr_tr_inv_k_tr_te = (
            l_k_tr_tr_inv_sigma_t_at_l_k_tr_tr_inv_k_tr_te.t()
            @ l_k_tr_tr_inv_sigma_t_at_l_k_tr_tr_inv_k_tr_te
        )

        # (n_te, n_te)
        predictive_cov = (
            k_te_te
            - k_tr_tr_inv_k_tr_te
            + compose_l_k_tr_tr_inv_sigma_t_at_l_k_tr_tr_inv_k_tr_te
        )

        predictive_cov += torch.exp(self.log_sigma) * torch.eye(
            k_te_te.shape[0], device=k_te_te.device
        )

        return predictive_mean, predictive_cov
//  [python]
def forward(self, x_te):
        """Forward pass.

        Parameters
        ----------
        x_te : `torch.Tensor`, `(x_te, hidden_dimension)`
            Test input.

        Returns
        -------
        predictive_mean : `torch.Tensor`, `(x_te, )`
            Predictive mean.

        predictive_cov : `torch.Tensor`, `(x_te, x_te)`
            Preditive covariance.

        """

        # get the kernels
        k_tr_tr, k_tr_te, k_te_tr, k_te_te = (
            self._k_tr_tr(),
            self._k_tr_te(x_te),
            self._k_te_tr(x_te),
            self._k_te_te(x_te),
        )

        # (n_tr, n_tr)
        l_k_tr_tr = torch.linalg.cholesky(k_tr_tr, upper=False)

        # (n_tr, 1)
        l_k_tr_tr_inv_mu, _ = torch.triangular_solve(
            input=self.y_tr_mu, A=l_k_tr_tr, upper=False
        )

        # (n_tr, n_tr)
        l_k_tr_tr_inv_sigma, _ = torch.triangular_solve(
            input=self._y_tr_sigma(), A=l_k_tr_tr, upper=False
        )

        # (n_tr, te)
        l_k_tr_tr_inv_k_tr_te, _ = torch.triangular_solve(
            input=k_tr_te, A=l_k_tr_tr, upper=False
        )

        # (n_te, 1)
        predictive_mean = l_k_tr_tr_inv_k_tr_te.t() @ l_k_tr_tr_inv_mu

        # (n_tr, n_te)
        k_tr_tr_inv_k_tr_te = (
            l_k_tr_tr_inv_k_tr_te.t() @ l_k_tr_tr_inv_k_tr_te
        )

        # (n_te, n_te)
        l_k_tr_tr_inv_sigma_t_at_l_k_tr_tr_inv_k_tr_te = (
            l_k_tr_tr_inv_sigma.t() @ l_k_tr_tr_inv_k_tr_te
        )

        # (n_te, n_te)
        compose_l_k_tr_tr_inv_sigma_t_at_l_k_tr_tr_inv_k_tr_te = (
            l_k_tr_tr_inv_sigma_t_at_l_k_tr_tr_inv_k_tr_te.t()
            @ l_k_tr_tr_inv_sigma_t_at_l_k_tr_tr_inv_k_tr_te
        )

        # (n_te, n_te)
        predictive_cov = (
            k_te_te
            - k_tr_tr_inv_k_tr_te
            + compose_l_k_tr_tr_inv_sigma_t_at_l_k_tr_tr_inv_k_tr_te
        )

        predictive_cov += torch.exp(self.log_sigma) * torch.eye(
            k_te_te.shape[0], device=k_te_te.device
        )

        return predictive_mean, predictive_cov
//  [python]
def _apply_floor_(self, A, B, component_batches):
        #B = self._apply_floor_scalar(B, self._max_abs_eig(B) * 1e-4)[0]  # To prevent Cholesky from failing
        L = torch.cholesky(B)
        L_inv = torch.inverse(L)
        num_floored = 0
        batch_size = component_batches[0][1] - component_batches[0][0]
        l = torch.zeros(batch_size, self.feat_dim, dtype=self.means.dtype, device=self.t_matrix.device)
        U = torch.zeros(batch_size, self.feat_dim, self.feat_dim, dtype=self.means.dtype, device=self.t_matrix.device) 
        for bstart, bend in component_batches:
            D = L_inv.matmul(A[bstart:bend, :, :]).matmul(L_inv.t())   
            for c in range(batch_size):
                l[c, :], U[c, :, :] = torch.symeig(D[c, :, :], eigenvectors=True)
            num_floored += torch.sum(l < 1).item()
            l = torch.clamp(l, min=1)
            D = U.matmul(l.unsqueeze(2) * U.transpose(1,2))
            A[bstart:bend, :, :] = L.matmul(D.transpose(1, 2)).matmul(L.t())
        print('Floored {:.1%} of the eigenvalues...'.format(num_floored / (self.n_components * self.feat_dim)))