//  [python]
def minibatch_KLD_upper_bound_iter(covar_module0, covar_module1, likelihood, latent_dim, m, H, train_xt, mu, log_v, z, P, P_in_current_batch, N, natural_gradient, id_covariate, eps):
    """
    Efficient unbiased estimate of the KL-divergence upper bound that enables the use of mini-batching.
    See L-VAE supplementary material.

    :param covar_module0: additive kernel (sum of cross-covariances) without id covariate
    :param covar_module1: additive kernel (sum of cross-covariances) with id covariate
    :param likelihood: GPyTorch likelihood model
    :param m: mean of inducing values u
    :param H: covariate matrix of inducing values u
    :param train_xt: auxiliary covariate information
    :param mu: variational mean
    :param log_v: (log) variational variance
    :param z: inducing points
    :param P: total number of subjects
    :param P_in_current_batch: number of subjects in the batch
    :param N: total number of samples
    :param eps: jitter
    :return: Unbiased estimate of the KL-divergence upper bound and its gradients w.r.t. m and H
    """

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    M = H.shape[-1]

    K0xz = covar_module0(train_xt, z).evaluate()
    K0zz = covar_module0(z, z).evaluate()
    K0zz = K0zz + eps * torch.eye(M, dtype=torch.double).to(device)
    LK0zz = torch.linalg.cholesky(K0zz)
    iK0zz = torch.linalg.solve(K0zz, torch.eye(M, dtype=torch.double).to(device))
    LH = torch.linalg.cholesky(H)
    iH = torch.linalg.solve(H, torch.eye(M, dtype=torch.double).to(device))

    A_part = (torch.matmul(torch.matmul(K0xz, iK0zz), m).squeeze() - mu.T).unsqueeze(dim=2)
    E_part = torch.matmul(torch.matmul(iK0zz, H), iK0zz)

    A = torch.tensor([0.0], dtype=torch.double).to(device)
    B = torch.tensor([0.0], dtype=torch.double).to(device)
    C = torch.tensor([0.0], dtype=torch.double).to(device)
    D = torch.tensor([0.0], dtype=torch.double).to(device)
    E = torch.tensor([0.0], dtype=torch.double).to(device)
    if natural_gradient:
        ng_P1 = torch.zeros(latent_dim, M, 1, dtype=torch.double).to(device)
        ng_P2 = torch.zeros(latent_dim, M, M, dtype=torch.double).to(device)

    subjects = torch.unique(train_xt[:, id_covariate]).tolist()
    for s in subjects:
        indices = train_xt[:, id_covariate] == s
        tx = train_xt[indices]
        T = tx.shape[0] 
        stacked_tx = torch.stack([tx for i in range(latent_dim)], dim=0)
        K0_st = covar_module0(stacked_tx, stacked_tx).evaluate()
        B_st = covar_module1(stacked_tx, stacked_tx).evaluate() + torch.eye(T, dtype=torch.double).to(device) * likelihood.noise_covar.noise.unsqueeze(dim=2)

        LB_st = torch.linalg.cholesky(B_st)
        iB_st = torch.linalg.solve(B_st, torch.eye(T, dtype=torch.double).to(device)).squeeze(dim=0)
        K0xz_st = K0xz[:, indices]
        K0zx_iB_K0xz = torch.einsum('bik,bij,bjl->bkl', K0xz_st, iB_st, K0xz_st)

        A = A + torch.einsum('bji,bjk,bkl->b', A_part[:, indices], iB_st, A_part[:, indices]).sum()
        B = B + torch.sum(torch.diagonal(iB_st, dim1=-1, dim2=-2).reshape(latent_dim, -1)*torch.exp(log_v[indices].T))
        C = C + 2 * torch.sum(torch.log(torch.diagonal(LB_st, dim1=-2, dim2=-1)))
        D = D + torch.sum(iB_st * K0_st) - torch.sum(K0zx_iB_K0xz * iK0zz)
        E = E + torch.sum(E_part * K0zx_iB_K0xz)

        if natural_gradient:
            mu_p = mu[indices].transpose(-1,-2).unsqueeze(dim=2)
            K0zx = K0xz_st.transpose(-1,-2)
            ng_P1 = ng_P1 + torch.matmul(K0zx, torch.matmul(iB_st, mu_p))
            ng_P2 = ng_P2 + K0zx_iB_K0xz
 
    F = torch.sum(log_v)

    #Compute kld_qu_pu
    tr1 = torch.sum(iK0zz * H.transpose(-1, -2))
    qf1 = torch.sum(m * torch.matmul(iK0zz, m))
    logdetK = 2 * torch.sum(torch.log(torch.diagonal(LK0zz, dim1=-1, dim2=-2)))
    logdetH = 2 * torch.sum(torch.log(torch.diagonal(LH, dim1=-1, dim2=-2)))
    kld_qu_pu = 0.5 * (tr1 + qf1 - latent_dim*M + logdetK - logdetH)

    kld_total = P/P_in_current_batch*0.5*(A + B + C + D + E - F) + kld_qu_pu - latent_dim*N/2

    grad_m, grad_H = None, None
    if natural_gradient:
        B = torch.matmul(iK0zz, torch.matmul(ng_P2, iK0zz)) + iK0zz
        grad_m = -torch.matmul(iK0zz, ng_P1) + torch.matmul(B, m)
        grad_H = 0.5*(-iH + B)

    return kld_total, grad_m, grad_H