jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def singular_cholesky(X):
    epsilon=1E-6
    I = t.eye(X.shape[-1], device=X.device, dtype=X.dtype)
    LdL = t.cholesky(X + epsilon*I)
    dL = t.triangular_solve(I, LdL)[0].transpose(-1, -2)
    t.allclose(I, LdL@dL.transpose(-1, -2))

    L = LdL - (epsilon/2)*dL

    return L
//  [python]
def singular_cholesky(X):
    epsilon=1E-6
    I = t.eye(X.shape[-1], device=X.device, dtype=X.dtype)
    LdL = t.linalg.cholesky(X + epsilon*I)
    dL = t.triangular_solve(I, LdL)[0].transpose(-1, -2)
    t.allclose(I, LdL@dL.transpose(-1, -2))

    L = LdL - (epsilon/2)*dL

    return L
//  [python]
def test_cholesky_solve_batched_broadcasting(self, device, dtype):
        from numpy.linalg import solve
        from torch.testing._internal.common_utils import random_symmetric_pd_matrix

        def run_test(A_dims, b_dims, upper):
            A_matrix_size = A_dims[-1]
            A_batch_dims = A_dims[:-2]
            A = random_symmetric_pd_matrix(A_matrix_size, *A_batch_dims,
                                           dtype=dtype, device='cpu')
            b = torch.randn(*b_dims, dtype=dtype, device='cpu')
            x_exp = torch.tensor(solve(A.numpy(), b.numpy()), dtype=dtype, device=device)
            A, b = A.to(dtype=dtype, device=device), b.to(dtype=dtype, device=device)
            L = torch.cholesky(A, upper)
            x = torch.cholesky_solve(b, L, upper=upper)
            self.assertEqual(x, x_exp)

        # test against numpy.linalg.solve
        for upper in [True, False]:
            run_test((2, 1, 3, 4, 4), (2, 1, 3, 4, 6), upper)  # no broadcasting
            run_test((2, 1, 3, 4, 4), (4, 6), upper)  # broadcasting b
            run_test((4, 4), (2, 1, 3, 4, 2), upper)  # broadcasting A
            run_test((1, 3, 1, 4, 4), (2, 1, 3, 4, 5), upper)
//  [python]
def test_cholesky_solve_batched_broadcasting(self, device, dtype):
        from numpy.linalg import solve
        from torch.testing._internal.common_utils import random_symmetric_pd_matrix

        def run_test(A_dims, b_dims, upper):
            A_matrix_size = A_dims[-1]
            A_batch_dims = A_dims[:-2]
            A = random_symmetric_pd_matrix(A_matrix_size, *A_batch_dims,
                                           dtype=dtype, device='cpu')
            b = torch.randn(*b_dims, dtype=dtype, device='cpu')
            x_exp = torch.tensor(solve(A.numpy(), b.numpy()), dtype=dtype, device=device)
            A, b = A.to(dtype=dtype, device=device), b.to(dtype=dtype, device=device)
            L = torch.linalg.cholesky(A, upper)
            x = torch.cholesky_solve(b, L, upper=upper)
            self.assertEqual(x, x_exp)

        # test against numpy.linalg.solve
        for upper in [True, False]:
            run_test((2, 1, 3, 4, 4), (2, 1, 3, 4, 6), upper)  # no broadcasting
            run_test((2, 1, 3, 4, 4), (4, 6), upper)  # broadcasting b
            run_test((4, 4), (2, 1, 3, 4, 2), upper)  # broadcasting A
            run_test((1, 3, 1, 4, 4), (2, 1, 3, 4, 5), upper)
//  [python]
def qmc_sample_fast(mu, cov, n_sample, rng=None):
    rng = torch.quasirandom.SobolEngine(dimension=2, scramble=True, seed=0)
    qr_seq = torch.stack([box_muller_transform(rng.draw(n_sample)) for _ in range(mu.size(0))], dim=1).unsqueeze(dim=2).type_as(mu)
    sample = mu + (torch.cholesky(cov) @ qr_seq.unsqueeze(dim=-1)).squeeze(dim=-1)
    return sample
//  [python]
def qmc_sample_fast(mu, cov, n_sample, rng=None):
    rng = torch.quasirandom.SobolEngine(dimension=2, scramble=True, seed=0)
    qr_seq = torch.stack([box_muller_transform(rng.draw(n_sample)) for _ in range(mu.size(0))], dim=1).unsqueeze(dim=2).type_as(mu)
    sample = mu + (torch.linalg.cholesky(cov) @ qr_seq.unsqueeze(dim=-1)).squeeze(dim=-1)
    return sample
//  [python]
def post_process(self, model, likelihood, train_loader):
        device = self.defaults['device']
        parameters = self.param_groups[0]['params']
        theta_star = parameters_to_vector(parameters).to(device)
        prior_prec = self.state['prior_prec']
        prior_mu = self.state['prior_mu']
        P = len(theta_star)
        JLJ = torch.zeros(P, P, device=device)
        G = torch.zeros(P, device=device)
        for data, target in train_loader:
            data, target = data.to(device), target.to(device)
            Js, Hess, rs = GGN(model, likelihood, data, target)
            JLJ += torch.einsum('mpk,mkl,mql->pq', Js, Hess, Js)
            G += torch.einsum('mpk,mk->p', Js, rs)
        # compute posterior covariance and precision
        self.state['precision'] = JLJ + prior_prec
        Chol = torch.cholesky(self.state['precision'])
        self.state['Sigma'] = torch.cholesky_inverse(Chol, upper=False)
        self.state['Sigma_chol'] = torch.cholesky(self.state['Sigma'], upper=False)
        # compute posterior mean according to BLR/Exact Laplace
        b = G + JLJ @ theta_star + prior_prec @ prior_mu
        self.state['mu'] = torch.cholesky_solve(b.reshape(-1, 1), Chol, upper=False).flatten()
        return self
//  [python]
def post_process(self, model, likelihood, train_loader):
        device = self.defaults['device']
        parameters = self.param_groups[0]['params']
        theta_star = parameters_to_vector(parameters).to(device)
        prior_prec = self.state['prior_prec']
        prior_mu = self.state['prior_mu']
        P = len(theta_star)
        JLJ = torch.zeros(P, P, device=device)
        G = torch.zeros(P, device=device)
        for data, target in train_loader:
            data, target = data.to(device), target.to(device)
            Js, Hess, rs = GGN(model, likelihood, data, target)
            JLJ += torch.einsum('mpk,mkl,mql->pq', Js, Hess, Js)
            G += torch.einsum('mpk,mk->p', Js, rs)
        # compute posterior covariance and precision
        self.state['precision'] = JLJ + prior_prec
        Chol = torch.linalg.cholesky(self.state['precision'])
        self.state['Sigma'] = torch.linalg.cholesky_inverse(Chol, upper=False)
        self.state['Sigma_chol'] = torch.linalg.cholesky(self.state['Sigma'], upper=False)
        # compute posterior mean according to BLR/Exact Laplace
        b = G + JLJ @ theta_star + prior_prec @ prior_mu
        self.state['mu'] = torch.linalg.solve_triangular(Chol, b.reshape(-1, 1), upper=False).flatten()
        return self
//  [python]
def forward(self, x):
        #x = self.relu(x)
        x = torch.cholesky(x)
        return x
//  [python]
def forward(self, x):
        #x = self.relu(x)
        x = torch.linalg.cholesky(x)
        return x
//  [python]
def isPD(A):
    try:
        torch.cholesky(A)
        return True
    except Exception:
        return False
//  [python]
def isPD(A):
    try:
        torch.linalg.cholesky(A)
        return True
    except Exception:
        return False
//  [python]
def test_dist_to_funsor_mvn(batch_shape, event_size):
    loc = torch.randn(batch_shape + (event_size,))
    cov = torch.randn(batch_shape + (event_size, 2 * event_size))
    cov = cov.matmul(cov.transpose(-1, -2))
    scale_tril = torch.cholesky(cov)
    d = dist.MultivariateNormal(loc, scale_tril=scale_tril)
    f = dist_to_funsor(d)
    assert isinstance(f, Funsor)

    value = d.sample()
    actual_log_prob = f(value=tensor_to_funsor(value, event_output=1))
    expected_log_prob = tensor_to_funsor(d.log_prob(value))
    assert_close(actual_log_prob, expected_log_prob)
//  [python]
def test_dist_to_funsor_mvn(batch_shape, event_size):
    loc = torch.randn(batch_shape + (event_size,))
    cov = torch.randn(batch_shape + (event_size, 2 * event_size))
    cov = cov.matmul(cov.transpose(-1, -2))
    scale_tril = torch.linalg.cholesky(cov)
    d = dist.MultivariateNormal(loc, scale_tril=scale_tril)
    f = dist_to_funsor(d)
    assert isinstance(f, Funsor)

    value = d.sample()
    actual_log_prob = f(value=tensor_to_funsor(value, event_output=1))
    expected_log_prob = tensor_to_funsor(d.log_prob(value))
    assert_close(actual_log_prob, expected_log_prob)
//  [python]
def scale_tril(self) -> torch.Tensor:
        return torch.cholesky(self.scale, upper=False)
//  [python]
def scale_tril(self) -> torch.Tensor:
        return torch.linalg.cholesky(self.scale, upper=False)
//  [python]
def get_KUKA_SGPMatrices_from_MDB(mbd:MultiRigidBody, dataset:StructTorchArray):
    ''' Given a KUKA model and configuration, get matrices needed for SGP and store into dataset
    '''
    # with torch.no_grad():

    # update model wit current position and velocity
    mbd.forwardKinematics(q=dataset.q , qDot=dataset.dq, qDDot=0*dataset.q)

    M, f, g = mbd.computationOfMfg()
    C = mbd.getJointFriction(dataset.dq)
    J_lambda, sigma_lambda = mbd.bilateralConstraints[0].getConstraintTerms()
    I_c = mbd.bilateralConstraints[0].I_c()
    A, b = J_lambda, -sigma_lambda
    I = torch.eye(mbd.nq, device=M.device)
    Ml = (A @ torch.solve(bT(A),M)[0]).inverse()
    L = torch.solve(bT(A),M)[0] @ Ml
    T = I - L @ A
    Minv = torch.solve(I,M)[0]
    ScFa = torch.cholesky(Minv)

    dataset.addColumns(
        M=M, f=f, g=g, C=C, A=A, b=b, Ml=Ml, L=L, T=T, Minv=Minv, ScFa=ScFa, I_c=I_c
    )
    if 'ddq' in dataset.keys():
        dataset.addColumns( ddqs = bmv(ScFa, dataset.ddq) )
    if 'ddqa' in dataset.keys():
        dataset.addColumns( ddqas = bmv(ScFa, dataset.ddqa) )
    
    return dataset
//  [python]
def get_KUKA_SGPMatrices_from_MDB(mbd:MultiRigidBody, dataset:StructTorchArray):
    ''' Given a KUKA model and configuration, get matrices needed for SGP and store into dataset
    '''
    # with torch.no_grad():

    # update model wit current position and velocity
    mbd.forwardKinematics(q=dataset.q , qDot=dataset.dq, qDDot=0*dataset.q)

    M, f, g = mbd.computationOfMfg()
    C = mbd.getJointFriction(dataset.dq)
    J_lambda, sigma_lambda = mbd.bilateralConstraints[0].getConstraintTerms()
    I_c = mbd.bilateralConstraints[0].I_c()
    A, b = J_lambda, -sigma_lambda
    I = torch.eye(mbd.nq, device=M.device)
    Ml = (A @ torch.solve(bT(A),M)[0]).inverse()
    L = torch.solve(bT(A),M)[0] @ Ml
    T = I - L @ A
    Minv = torch.solve(I,M)[0]
    ScFa = torch.linalg.cholesky(Minv)

    dataset.addColumns(
        M=M, f=f, g=g, C=C, A=A, b=b, Ml=Ml, L=L, T=T, Minv=Minv, ScFa=ScFa, I_c=I_c
    )
    if 'ddq' in dataset.keys():
        dataset.addColumns( ddqs = bmv(ScFa, dataset.ddq) )
    if 'ddqa' in dataset.keys():
        dataset.addColumns( ddqas = bmv(ScFa, dataset.ddqa) )
    
    return dataset
//  [python]
def survival_likelihood(self, X, X_f, y, y_f, x_cat, x_cat_f, mask):
        x_concat_S = torch.cat([X[~mask],y[~mask]],dim=1)
        cov_mat_S = self.r(x_concat_S)
        L_S = torch.cholesky(cov_mat_S) @ torch.randn_like(y[~mask])
        S, h_S = self.m_q.forward_S(X, y, mask, x_cat,L_S)
        if S.numel() == 0:
            S = S.detach()
        y_f = torch.autograd.Variable(y_f, requires_grad=True)
        x_concat_f = torch.cat([X_f,y_f],dim=1)
        cov_mat_f = self.r(x_concat_f)
        L_f = torch.linalg.cholesky(cov_mat_f)@ torch.randn_like(y_f)
        f, h_f = self.m_q.forward_f(X_f, y_f, x_cat_f,L_f)
        if f.numel() == 0:
            f = f.detach()
        total_loss = log_objective_mean(S, f)
        tmp = torch.ones_like(y) * self.m_p
        reg = torch.mean((torch.cat([h_f,h_S],dim=0) - tmp) ** 2)
        return total_loss,reg
//  [python]
def survival_likelihood(self, X, X_f, y, y_f, x_cat, x_cat_f, mask):
        x_concat_S = torch.cat([X[~mask],y[~mask]],dim=1)
        cov_mat_S = self.r(x_concat_S)
        L_S = torch.linalg.cholesky(cov_mat_S) @ torch.randn_like(y[~mask])
        S, h_S = self.m_q.forward_S(X, y, mask, x_cat,L_S)
        if S.numel() == 0:
            S = S.detach()
        y_f = torch.autograd.Variable(y_f, requires_grad=True)
        x_concat_f = torch.cat([X_f,y_f],dim=1)
        cov_mat_f = self.r(x_concat_f)
        L_f = torch.linalg.cholesky(cov_mat_f)@ torch.randn_like(y_f)
        f, h_f = self.m_q.forward_f(X_f, y_f, x_cat_f,L_f)
        if f.numel() == 0:
            f = f.detach()
        total_loss = log_objective_mean(S, f)
        tmp = torch.ones_like(y) * self.m_p
        reg = torch.mean((torch.cat([h_f,h_S],dim=0) - tmp) ** 2)
        return total_loss,reg
//  [python]
def cholesky(X):
    """Cholesky decomposition of $X$.

    >>> cholesky(X)
    return L
    torch.mm(L, L.t()) = X
    """

    if torch.__version__ >= '1.0.0':
        return torch.cholesky(X)
    else:
        return torch.potrf(X, upper=False)
//  [python]
def cholesky(X):
    """Cholesky decomposition of $X$.

    >>> cholesky(X)
    return L
    torch.mm(L, L.t()) = X
    """

    if torch.__version__ >= '1.0.0':
        return torch.linalg.cholesky(X)
    else:
        return torch.potrf(X, upper=False)
//  [python]
def _setup_task_dicts(self, train_data_tuples):
        task_dicts, parameters = [], []

        for train_x, train_y in train_data_tuples:
            task_dict = OrderedDict()

            # a) prepare data
            x_tensor, y_tensor = self._prepare_data_per_task(train_x, train_y)
            task_dict['train_x'], task_dict['train_y'] = x_tensor, y_tensor

            prior_param_sample = self.hyper_posterior.sample(sample_shape=(20,))
            mean_module, covar_module = self._aggregate_gp_priors(prior_param_sample)

            task_dict['gp_model'] = LearnedGPRegressionModelApproximate(x_tensor, y_tensor, self.likelihood,
                                                                        mean_module=mean_module,
                                                                        covar_module=covar_module)
            init_mean = mean_module(x_tensor)
            init_mean += 1e-3 * torch.randn_like(init_mean)
            task_dict['gp_model'].variational_distribution.variational_mean.data.copy_(init_mean)

            prior_covar = covar_module(x_tensor)
            init_chol_covar = torch.cholesky(prior_covar + 1e-3 * torch.eye(prior_covar.shape[0]))
            task_dict['gp_model'].variational_distribution.chol_variational_covar.data.copy_(init_chol_covar)

            parameters.extend(task_dict['gp_model'].variational_parameters())
            task_dicts.append(task_dict)

        return task_dicts, parameters
//  [python]
def _setup_task_dicts(self, train_data_tuples):
        task_dicts, parameters = [], []

        for train_x, train_y in train_data_tuples:
            task_dict = OrderedDict()

            # a) prepare data
            x_tensor, y_tensor = self._prepare_data_per_task(train_x, train_y)
            task_dict['train_x'], task_dict['train_y'] = x_tensor, y_tensor

            prior_param_sample = self.hyper_posterior.sample(sample_shape=(20,))
            mean_module, covar_module = self._aggregate_gp_priors(prior_param_sample)

            task_dict['gp_model'] = LearnedGPRegressionModelApproximate(x_tensor, y_tensor, self.likelihood,
                                                                        mean_module=mean_module,
                                                                        covar_module=covar_module)
            init_mean = mean_module(x_tensor)
            init_mean += 1e-3 * torch.randn_like(init_mean)
            task_dict['gp_model'].variational_distribution.variational_mean.data.copy_(init_mean)

            prior_covar = covar_module(x_tensor)
            init_chol_covar = torch.linalg.cholesky(prior_covar + 1e-3 * torch.eye(prior_covar.shape[0]))
            task_dict['gp_model'].variational_distribution.chol_variational_covar.data.copy_(init_chol_covar)

            parameters.extend(task_dict['gp_model'].variational_parameters())
            task_dicts.append(task_dict)

        return task_dicts, parameters
//  [python]
def cholesky_update(self, hyper):
        self.gram_mat_update(hyper)
        eye_mat = torch.diag(self.gram_mat.new_ones(self.gram_mat.size(0)))
        for jitter_const in [0, 1e-8, 1e-7, 1e-6, 1e-5, 1e-4, 1e-3]:
            chol_jitter = torch.trace(self.gram_mat).item() * jitter_const
            try:
                # cholesky is lower triangular matrix
                self.cholesky = torch.cholesky(self.gram_mat + eye_mat * chol_jitter, upper=False)
                self.jitter = chol_jitter
                return
            except RuntimeError:
                pass
        raise RuntimeError('Absolute entry values of Gram matrix are between %.4E~%.4E with trace %.4E' %
                           (torch.min(torch.abs(self.gram_mat)).item(), torch.max(torch.abs(self.gram_mat)).item(),
                            torch.trace(self.gram_mat).item()))
//  [python]
def cholesky_update(self, hyper):
        self.gram_mat_update(hyper)
        eye_mat = torch.diag(self.gram_mat.new_ones(self.gram_mat.size(0)))
        for jitter_const in [0, 1e-8, 1e-7, 1e-6, 1e-5, 1e-4, 1e-3]:
            chol_jitter = torch.trace(self.gram_mat).item() * jitter_const
            try:
                # cholesky is lower triangular matrix
                self.cholesky = torch.linalg.cholesky(self.gram_mat + eye_mat * chol_jitter, upper=False)
                self.jitter = chol_jitter
                return
            except RuntimeError:
                pass
        raise RuntimeError('Absolute entry values of Gram matrix are between %.4E~%.4E with trace %.4E' %
                           (torch.min(torch.abs(self.gram_mat)).item(), torch.max(torch.abs(self.gram_mat)).item(),
                            torch.trace(self.gram_mat).item()))
//  [python]
def cholesky_wrapper(mat, default_diag=None, force_cpu=False):
    device = mat.device
    if force_cpu:
        mat = mat.cpu()
    try:
        tril = torch.cholesky(mat, upper=False)
    except RuntimeError:
        n_dims = mat.size(-1)
        tril = []
        default_tril_single = torch.diag(mat.new_tensor(default_diag)) if default_diag is not None \
            else torch.eye(n_dims, dtype=mat.dtype, device=mat.device)
        for cov in mat.reshape(-1, n_dims, n_dims):
            try:
                tril.append(torch.cholesky(cov, upper=False))
            except RuntimeError:
                tril.append(default_tril_single)
        tril = torch.stack(tril, dim=0).reshape(mat.shape)
    return tril.to(device)
//  [python]
def cholesky_wrapper(mat, default_diag=None, force_cpu=False):
    device = mat.device
    if force_cpu:
        mat = mat.cpu()
    try:
        tril = torch.linalg.cholesky(mat, upper=False)
    except RuntimeError:
        n_dims = mat.size(-1)
        tril = []
        default_tril_single = torch.diag(mat.new_tensor(default_diag)) if default_diag is not None \
            else torch.eye(n_dims, dtype=mat.dtype, device=mat.device)
        for cov in mat.reshape(-1, n_dims, n_dims):
            try:
                tril.append(torch.linalg.cholesky(cov, upper=False))
            except RuntimeError:
                tril.append(default_tril_single)
        tril = torch.stack(tril, dim=0).reshape(mat.shape)
    return tril.to(device)
//  [python]
def predict_f_samples(self, Xnew, num_samples):
        """
        Produce samples from the posterior latent function(s) at the points
        Xnew.
        """
        mu, var = self.predict_f(Xnew, full_cov=True)
        jitter = torch.eye(mu.size(0), dtype=mu.dtype, device=mu.device) * self.jitter_level
        samples = []
        for i in range(self.num_latent): # TV-Todo: batch??
            L = torch.cholesky(var[:, :, i] + jitter, upper=False)
            V = torch.randn(L.size(0), num_samples, dtype=L.dtype, device=L.device)
            samples.append(mu[:, i:i + 1] + torch.matmul(L, V))
        return torch.stack(samples, dim=0)
//  [python]
def predict_f_samples(self, Xnew, num_samples):
        """
        Produce samples from the posterior latent function(s) at the points
        Xnew.
        """
        mu, var = self.predict_f(Xnew, full_cov=True)
        jitter = torch.eye(mu.size(0), dtype=mu.dtype, device=mu.device) * self.jitter_level
        samples = []
        for i in range(self.num_latent): # TV-Todo: batch??
            L = torch.linalg.cholesky(var[:, :, i] + jitter, upper=False)
            V = torch.randn(L.size(0), num_samples, dtype=L.dtype, device=L.device)
            samples.append(mu[:, i:i + 1] + torch.matmul(L, V))
        return torch.stack(samples, dim=0)
//  [python]
def forward(ctx, H, b):
        chol = torch.cholesky(H)
        xx = torch.cholesky_solve(b, chol)
        ctx.save_for_backward(chol, xx)

        return xx
//  [python]
def forward(ctx, H, b):
        chol = torch.linalg.cholesky(H)
        xx = torch.cholesky_solve(b, chol)
        ctx.save_for_backward(chol, xx)

        return xx
//  [python]
def __init__(self, loc, scale, distance_matrix, lengthscale):
        cov = torch.exp(- distance_matrix / lengthscale) * scale ** 2.0
        chol = torch.cholesky(cov)
        super().__init__(loc=loc, scale_tril=chol)
//  [python]
def __init__(self, loc, scale, distance_matrix, lengthscale):
        cov = torch.exp(- distance_matrix / lengthscale) * scale ** 2.0
        chol = torch.linalg.cholesky(cov)
        super().__init__(loc=loc, scale_tril=chol)
//  [python]
def estimate_variance(self, var0, invert=True):
        tau = 1/var0

        U, V = self.hessians

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

            if name == 'bias':
                continue

            U_ = U[(module, name)].clone()
            V_ = V[(module, name)].clone()

            if invert:
                m, n = int(U_.shape[0]), int(V_.shape[0])

                U_ += torch.sqrt(tau)*torch.eye(m, device='cuda')
                V_ += torch.sqrt(tau)*torch.eye(n, device='cuda')

                # For numerical stability
                u = torch.cholesky(U_.cpu() + 1e-6*torch.eye(m))
                v = torch.cholesky(V_.cpu() + 1e-6*torch.eye(n))

                U_ = torch.cholesky(torch.cholesky_inverse(u), upper=False).cuda()
                V_ = torch.cholesky(torch.cholesky_inverse(v), upper=True).cuda()

            module.__getattr__(f'{name}_U_half').copy_(U_)
            module.__getattr__(f'{name}_V_half').copy_(V_)
//  [python]
def estimate_variance(self, var0, invert=True):
        tau = 1/var0

        U, V = self.hessians

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

            if name == 'bias':
                continue

            U_ = U[(module, name)].clone()
            V_ = V[(module, name)].clone()

            if invert:
                m, n = int(U_.shape[0]), int(V_.shape[0])

                U_ += torch.sqrt(tau)*torch.eye(m, device='cuda')
                V_ += torch.sqrt(tau)*torch.eye(n, device='cuda')

                # For numerical stability
                u = torch.linalg.cholesky(U_.cpu() + 1e-6*torch.eye(m))
                v = torch.linalg.cholesky(V_.cpu() + 1e-6*torch.eye(n))

                U_ = torch.linalg.cholesky_ex(u, upper=False).inverse().L.cuda()
                V_ = torch.linalg.cholesky_ex(v, upper=True).inverse().L.cuda()

            module.__getattr__(f'{name}_U_half').copy_(U_)
            module.__getattr__(f'{name}_V_half').copy_(V_)
//  [python]
def __init__(self, units=64):
        """Initializes a dynamics model for our door task."""

        super().__init__(state_dim=2)

        # Fixed dynamics covariance
        self.Q_scale_tril = nn.Parameter(
            torch.cholesky(torch.diag(torch.FloatTensor([0.02, 0.02]))),
            requires_grad=False,
        )

        # Build neural network
        self.state_layers = layers.state_layers(units=units)
        self.control_layers = layers.control_layers(units=units)
        self.shared_layers = nn.Sequential(
            nn.Linear(units * 2, units),
            resblocks.Linear(units),
            resblocks.Linear(units),
            resblocks.Linear(units),
            nn.Linear(units, self.state_dim + 1),
        )
        self.units = units
//  [python]
def __init__(self, units=64):
        """Initializes a dynamics model for our door task."""

        super().__init__(state_dim=2)

        # Fixed dynamics covariance
        self.Q_scale_tril = nn.Parameter(
            torch.linalg.cholesky(torch.diag(torch.FloatTensor([0.02, 0.02]))),
            requires_grad=False,
        )

        # Build neural network
        self.state_layers = layers.state_layers(units=units)
        self.control_layers = layers.control_layers(units=units)
        self.shared_layers = nn.Sequential(
            nn.Linear(units * 2, units),
            resblocks.Linear(units),
            resblocks.Linear(units),
            resblocks.Linear(units),
            nn.Linear(units, self.state_dim + 1),
        )
        self.units = units
//  [python]
def loss(self, *args):
        # https: // www.uta.edu / math / _docs / preprint / 2014 / rep2014_04.pdf
        # H is n_views * n_samples * k

        # Subtract the mean from each output
        views = [view - view.mean(dim=0) for view in args]

        # Concatenate all views and from this get the cross-covariance matrix
        all_views = torch.cat(views, dim=1)
        C = torch.matmul(all_views.T, all_views)

        # Get the block covariance matrix placing Xi^TX_i on the diagonal
        D = block_diag([torch.matmul(view.T, view) for view in views])

        # In MCCA our eigenvalue problem Cv = lambda Dv
        C = C - D

        # Use the cholesky method to whiten the matrix C R^{-1}CRv = lambda v
        R = torch.cholesky(D, upper=True)

        C_whitened = torch.inverse(R.T) @ C @ torch.inverse(R)

        [eigvals, eigvecs] = torch.symeig(C_whitened, eigenvectors=True)

        # Sort eigenvalues so largest first
        idx = torch.argsort(eigvals, descending=True)

        # Sum the first #outdim_size values (after subtracting 1).
        corr = (eigvals[idx][:self.outdim_size] - 1).sum()

        return -corr
//  [python]
def loss(self, *args):
        # https: // www.uta.edu / math / _docs / preprint / 2014 / rep2014_04.pdf
        # H is n_views * n_samples * k

        # Subtract the mean from each output
        views = [view - view.mean(dim=0) for view in args]

        # Concatenate all views and from this get the cross-covariance matrix
        all_views = torch.cat(views, dim=1)
        C = torch.matmul(all_views.T, all_views)

        # Get the block covariance matrix placing Xi^TX_i on the diagonal
        D = block_diag([torch.matmul(view.T, view) for view in views])

        # In MCCA our eigenvalue problem Cv = lambda Dv
        C = C - D

        # Use the cholesky method to whiten the matrix C R^{-1}CRv = lambda v
        R = torch.linalg.cholesky(D, upper=True)

        C_whitened = torch.inverse(R.T) @ C @ torch.inverse(R)

        [eigvals, eigvecs] = torch.symeig(C_whitened, eigenvectors=True)

        # Sort eigenvalues so largest first
        idx = torch.argsort(eigvals, descending=True)

        # Sum the first #outdim_size values (after subtracting 1).
        corr = (eigvals[idx][:self.outdim_size] - 1).sum()

        return -corr
//  [python]
def _sample_posterior(self, x, num_samples, context=None):
        log_weights = torch.log(self.module.soft_max(self.module.soft_weights))
        T = self.module.covars[None, :, :, :] + x[1][:, None, :, :]

        p_weights = log_weights + dist.MultivariateNormal(
            loc=self.module.means, covariance_matrix=T
        ).log_prob(x[0][:, None, :])
        p_weights -= torch.logsumexp(p_weights, axis=1)[:, None]

        L_t = torch.cholesky(T)
        T_inv = torch.cholesky_solve(
            torch.eye(self.d, device=self.device), L_t)

        diff = x[0][:, None, :] - self.module.means
        T_prod = torch.matmul(T_inv, diff[:, :, :, None])
        p_means = self.module.means + torch.matmul(
            self.module.covars,
            T_prod
        ).squeeze()

        p_covars = self.module.covars - torch.matmul(
            self.module.covars,
            torch.matmul(T_inv, self.module.covars)
        )

        idx = dist.Categorical(logits=p_weights).sample([num_samples])
        samples = dist.MultivariateNormal(
            loc=p_means, covariance_matrix=p_covars).sample([num_samples])

        return samples.transpose(0, 1)[
            torch.arange(len(x), device=self.device)[:, None, None, None],
            torch.arange(num_samples, device=self.device)[None, :, None, None],
            idx.T[:, :, None, None],
            torch.arange(self.d, device=self.device)[None, None, None, :]
        ].squeeze()
//  [python]
def _sample_posterior(self, x, num_samples, context=None):
        log_weights = torch.log(self.module.soft_max(self.module.soft_weights))
        T = self.module.covars[None, :, :, :] + x[1][:, None, :, :]

        p_weights = log_weights + dist.MultivariateNormal(
            loc=self.module.means, covariance_matrix=T
        ).log_prob(x[0][:, None, :])
        p_weights -= torch.logsumexp(p_weights, axis=1)[:, None]

        L_t = torch.linalg.cholesky(T)
        T_inv = torch.linalg.solve(
            torch.eye(self.d, device=self.device), L_t)

        diff = x[0][:, None, :] - self.module.means
        T_prod = torch.matmul(T_inv, diff[:, :, :, None])
        p_means = self.module.means + torch.matmul(
            self.module.covars,
            T_prod
        ).squeeze()

        p_covars = self.module.covars - torch.matmul(
            self.module.covars,
            torch.matmul(T_inv, self.module.covars)
        )

        idx = dist.Categorical(logits=p_weights).sample([num_samples])
        samples = dist.MultivariateNormal(
            loc=p_means, covariance_matrix=p_covars).sample([num_samples])

        return samples.transpose(0, 1)[
            torch.arange(len(x), device=self.device)[:, None, None, None],
            torch.arange(num_samples, device=self.device)[None, :, None, None],
            idx.T[:, :, None, None],
            torch.arange(self.d, device=self.device)[None, None, None, :]
        ].squeeze()
//  [python]
def solve(self, X, y, lmbd, X_test=None, Y_test=None, indices=None):
        if indices is None:
            X_transposed_X = torch.matmul(torch.t(X), X) + lmbd * torch.eye(X.size()[1], device=X.device, dtype=torch.float64)
            R = torch.cholesky(X_transposed_X)
        to_return = {}
        X_original = X.clone()
        y_original = y.clone()
        for i in range(0, 4):
            if indices is not None:
                X = X_original[indices[i]]
                y = y_original[indices[i]]
                X_transposed_X = torch.matmul(torch.t(X), X) + lmbd * torch.eye(X.size()[1], device=X.device, dtype=torch.float64)
                R = torch.cholesky(X_transposed_X)
            y_torch_i = y[:, i]
            z = torch.triangular_solve(torch.matmul(torch.t(X), y_torch_i).view(X.size()[1], 1), R, upper=False).solution
            w = torch.triangular_solve(z, torch.t(R)).solution.view(X.size()[1])
            losses = 0.5 * torch.pow((torch.matmul(X, w) - y_torch_i), 2)
            to_return[str(i)] = {'weights': w.to('cuda').type(torch.float32),
                                 'losses': losses.type(torch.float32)}
        return to_return
//  [python]
def solve(self, X, y, lmbd, X_test=None, Y_test=None, indices=None):
        if indices is None:
            X_transposed_X = torch.matmul(torch.t(X), X) + lmbd * torch.eye(X.size()[1], device=X.device, dtype=torch.float64)
            R = torch.linalg.cholesky(X_transposed_X)
        to_return = {}
        X_original = X.clone()
        y_original = y.clone()
        for i in range(0, 4):
            if indices is not None:
                X = X_original[indices[i]]
                y = y_original[indices[i]]
                X_transposed_X = torch.matmul(torch.t(X), X) + lmbd * torch.eye(X.size()[1], device=X.device, dtype=torch.float64)
                R = torch.linalg.cholesky(X_transposed_X)
            y_torch_i = y[:, i]
            z = torch.triangular_solve(torch.matmul(torch.t(X), y_torch_i).view(X.size()[1], 1), R, upper=False).solution
            w = torch.triangular_solve(z, torch.t(R)).solution.view(X.size()[1])
            losses = 0.5 * torch.pow((torch.matmul(X, w) - y_torch_i), 2)
            to_return[str(i)] = {'weights': w.to('cuda').type(torch.float32),
                                 'losses': losses.type(torch.float32)}
        return to_return
//  [python]
def forward(self, observations):
        N, observation_dim = observations.shape
        assert self.observation_dim == observation_dim
        predicted_state = (self.H_pinv[None, :, :] @ observations[:, :, None]).squeeze(-1)
        # Lower-triangular cholesky decomposition of covariance
        covariance = torch.cholesky(self.H_pinv @ self.R @ self.R.transpose(-1, -2) @ self.H_pinv.transpose(-1, -2))
        
        return predicted_state, covariance.expand((N, self.state_dim, self.state_dim))
//  [python]
def forward(self, observations):
        N, observation_dim = observations.shape
        assert self.observation_dim == observation_dim
        predicted_state = (self.H_pinv[None, :, :] @ observations[:, :, None]).squeeze(-1)
        # Lower-triangular cholesky decomposition of covariance
        covariance = torch.linalg.cholesky(self.H_pinv @ self.R @ self.R.transpose(-1, -2) @ self.H_pinv.transpose(-1, -2))
        
        return predicted_state, covariance.expand((N, self.state_dim, self.state_dim))
//  [python]
def predict(covar_module0, covar_module1, likelihood, train_xt, test_x, mu, z, P, T, id_covariate, eps):
    """
    Helper function to perform predictions.
    
    """
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    Q = train_xt.shape[1]
    M = z.shape[0]
    I_M = torch.eye(M, dtype=torch.double).to(device)
    I_T = torch.eye(T, dtype=torch.double).to(device)

    x_st = torch.reshape(train_xt, [P, T, Q])
    mu_st = torch.reshape(mu, [P, T, 1])

    K0xz = covar_module0(train_xt, z).evaluate()
    K0zz = covar_module0(z, z).evaluate() + eps * I_M
    K1_st = covar_module1(x_st, x_st).evaluate()
    K0Xz = covar_module0(test_x, z).evaluate()

    B_st = K1_st + I_T * likelihood.noise_covar.noise
    LB_st = torch.cholesky(B_st)
    iB_st = torch.cholesky_solve(I_T, LB_st)
    K0xz_st = torch.reshape(K0xz, [P, T, M])
    K0zx_st = K0xz_st.transpose(-1, -2)

    iB_K0xz = torch.matmul(iB_st, K0xz_st)
    K0zx_iB_K0xz = torch.matmul(K0xz.T, torch.reshape(iB_K0xz, [P*T, M]))
    H = K0zz + K0zx_iB_K0xz

    iB_mu = torch.matmul(iB_st, mu_st).view(-1)
    K0xz_iH_K0zx_iB_mu_st = torch.matmul(K0xz, torch.solve(torch.matmul(K0xz.T, iB_mu).unsqueeze(dim=1), H)[0]).reshape(P, T, -1)
    iB_K0xz_iH_K0zx_iB_mu = torch.matmul(iB_st, K0xz_iH_K0zx_iB_mu_st).view(-1)
    mu_tilde = iB_mu - iB_K0xz_iH_K0zx_iB_mu
    K0Xz_iK0zz_K0zx_mu_tilde = torch.matmul(K0Xz, torch.solve(torch.matmul(K0xz.T, mu_tilde).unsqueeze(dim=1), K0zz)[0]).squeeze()

    test_subjects = torch.unique(test_x[:, id_covariate]).cpu().numpy()
    mask = np.isin(train_xt[:, id_covariate].cpu().numpy(), test_subjects)

    K1Xx_mu_tilde = torch.zeros(test_x.shape[0], dtype=torch.double).to(device)
    for s in test_subjects:
        indices = test_x[:, id_covariate] == s
        K1Xx = covar_module1(test_x[indices], train_xt[mask]).evaluate()
        K1Xx_mu_tilde[indices] = torch.matmul(K1Xx, mu_tilde[mask])

    Z_pred = K0Xz_iK0zz_K0zx_mu_tilde + K1Xx_mu_tilde

    return Z_pred
//  [python]
def predict(covar_module0, covar_module1, likelihood, train_xt, test_x, mu, z, P, T, id_covariate, eps):
    """
    Helper function to perform predictions.
    
    """
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    Q = train_xt.shape[1]
    M = z.shape[0]
    I_M = torch.eye(M, dtype=torch.double).to(device)
    I_T = torch.eye(T, dtype=torch.double).to(device)

    x_st = torch.reshape(train_xt, [P, T, Q])
    mu_st = torch.reshape(mu, [P, T, 1])

    K0xz = covar_module0(train_xt, z).evaluate()
    K0zz = covar_module0(z, z).evaluate() + eps * I_M
    K1_st = covar_module1(x_st, x_st).evaluate()
    K0Xz = covar_module0(test_x, z).evaluate()

    B_st = K1_st + I_T * likelihood.noise_covar.noise
    LB_st = torch.linalg.cholesky(B_st)
    iB_st = torch.linalg.solve_triangular(LB_st, I_T)
    K0xz_st = torch.reshape(K0xz, [P, T, M])
    K0zx_st = K0xz_st.transpose(-1, -2)

    iB_K0xz = torch.matmul(iB_st, K0xz_st)
    K0zx_iB_K0xz = torch.matmul(K0xz.T, torch.reshape(iB_K0xz, [P*T, M]))
    H = K0zz + K0zx_iB_K0xz

    iB_mu = torch.matmul(iB_st, mu_st).view(-1)
    K0xz_iH_K0zx_iB_mu_st = torch.matmul(K0xz, torch.linalg.solve(H, torch.matmul(K0xz.T, iB_mu).unsqueeze(dim=1))).reshape(P, T, -1)
    iB_K0xz_iH_K0zx_iB_mu = torch.matmul(iB_st, K0xz_iH_K0zx_iB_mu_st).view(-1)
    mu_tilde = iB_mu - iB_K0xz_iH_K0zx_iB_mu
    K0Xz_iK0zz_K0zx_mu_tilde = torch.matmul(K0Xz, torch.linalg.solve(K0zz, torch.matmul(K0xz.T, mu_tilde).unsqueeze(dim=1))).squeeze()

    test_subjects = torch.unique(test_x[:, id_covariate]).cpu().numpy()
    mask = np.isin(train_xt[:, id_covariate].cpu().numpy(), test_subjects)

    K1Xx_mu_tilde = torch.zeros(test_x.shape[0], dtype=torch.double).to(device)
    for s in test_subjects:
        indices = test_x[:, id_covariate] == s
        K1Xx = covar_module1(test_x[indices], train_xt[mask]).evaluate()
        K1Xx_mu_tilde[indices] = torch.matmul(K1Xx, mu_tilde[mask])

    Z_pred = K0Xz_iK0zz_K0zx_mu_tilde + K1Xx_mu_tilde

    return Z_pred
//  [python]
def get_loss(self, loc, mu, cov, gt):
        loc_norm = box_muller_transform(loc).permute(2, 0, 1, 3).expand((loc.size(2),) + mu.shape)
        p_sample = mu + (torch.cholesky(cov) @ loc_norm.unsqueeze(dim=-1)).squeeze(dim=-1)

        # Distance loss
        loss_dist = (p_sample.mean(dim=1) - gt.permute(0, 3, 1, 2).mean(dim=1)).norm(p=2, dim=-1).min(dim=0)[0]
        loss_dist = loss_dist.mean()

        # Discrepancy loss
        loss_disc = (loc.unsqueeze(dim=2) - loc.unsqueeze(dim=3)).norm(p=2, dim=-1)
        loss_disc = loss_disc.topk(k=2, dim=-1, largest=False, sorted=True)[0][..., 1].log().mul(-1).mean(dim=-1)
        loss_disc = loss_disc.mean()

        return loss_dist, loss_disc
//  [python]
def get_loss(self, loc, mu, cov, gt):
        loc_norm = box_muller_transform(loc).permute(2, 0, 1, 3).expand((loc.size(2),) + mu.shape)
        p_sample = mu + (torch.linalg.cholesky(cov) @ loc_norm.unsqueeze(dim=-1)).squeeze(dim=-1)

        # Distance loss
        loss_dist = (p_sample.mean(dim=1) - gt.permute(0, 3, 1, 2).mean(dim=1)).norm(p=2, dim=-1).min(dim=0)[0]
        loss_dist = loss_dist.mean()

        # Discrepancy loss
        loss_disc = (loc.unsqueeze(dim=2) - loc.unsqueeze(dim=3)).norm(p=2, dim=-1)
        loss_disc = loss_disc.topk(k=2, dim=-1, largest=False, sorted=True)[0][..., 1].log().mul(-1).mean(dim=-1)
        loss_disc = loss_disc.mean()

        return loss_dist, loss_disc
//  [python]
def forward(self, observations):
        N, observation_dim = observations.shape
        assert self.observation_dim == observation_dim
        
        def expected_observations(x):
            return torch.linalg.norm(self.satXYZb[:, :3] - x[None, :3], dim=-1) + torch.abs(self.satXYZb[:, 3] - x[None, 3])
        
        predicted_state = torch.zeros(N, self.state_dim)
        covariance = torch.eye(self.state_dim).expand((N, self.state_dim, self.state_dim))
        
        for bidx, observation in enumerate(observations):
            initial = torch.zeros(4, requires_grad = True)
            for i in range(self.iterations): 
                previous_data = initial.clone()
                expec = expected_observations(initial)
                res = observation - expec
                jac = torch.autograd.functional.jacobian(expected_observations, initial.clone())
                # update 
                initial.data += (torch.linalg.pinv(jac) @ res[:, None]).squeeze(-1)
                # zero out current gradient to hold new gradients in next iteration 
                if torch.sum(torch.abs(initial - previous_data)) < self.convergence:
                    break
            predicted_state[bidx, self.idx_pos_mask] = initial[:3]
            predicted_state[bidx, self.idx_b_mask] = initial[3]
            # Lower-triangular cholesky decomposition of covariance
            R = torch.eye(self.observation_dim) * self.pr_std
            H = torch.autograd.functional.jacobian(expected_observations, predicted_state[bidx, :].clone())
            H_pinv = torch.linalg.pinv(H)
            covariance[bidx, :] = torch.cholesky(H_pinv @ R @ R.transpose(-1, -2) @ H_pinv.transpose(-1, -2) + torch.eye(self.state_dim)*1e-5)
        
        return predicted_state, covariance
//  [python]
def forward(self, observations):
        N, observation_dim = observations.shape
        assert self.observation_dim == observation_dim
        
        def expected_observations(x):
            return torch.linalg.norm(self.satXYZb[:, :3] - x[None, :3], dim=-1) + torch.abs(self.satXYZb[:, 3] - x[None, 3])
        
        predicted_state = torch.zeros(N, self.state_dim)
        covariance = torch.eye(self.state_dim).expand((N, self.state_dim, self.state_dim))
        
        for bidx, observation in enumerate(observations):
            initial = torch.zeros(4, requires_grad = True)
            for i in range(self.iterations): 
                previous_data = initial.clone()
                expec = expected_observations(initial)
                res = observation - expec
                jac = torch.autograd.functional.jacobian(expected_observations, initial.clone())
                # update 
                initial.data += (torch.linalg.pinv(jac) @ res[:, None]).squeeze(-1)
                # zero out current gradient to hold new gradients in next iteration 
                if torch.sum(torch.abs(initial - previous_data)) < self.convergence:
                    break
            predicted_state[bidx, self.idx_pos_mask] = initial[:3]
            predicted_state[bidx, self.idx_b_mask] = initial[3]
            # Lower-triangular cholesky decomposition of covariance
            R = torch.eye(self.observation_dim) * self.pr_std
            H = torch.autograd.functional.jacobian(expected_observations, predicted_state[bidx, :].clone())
            H_pinv = torch.linalg.pinv(H)
            covariance[bidx, :] = torch.linalg.cholesky(H_pinv @ R @ R.transpose(-1, -2) @ H_pinv.transpose(-1, -2) + torch.eye(self.state_dim)*1e-5)
        
        return predicted_state, covariance
//  [python]
def predict_decomposition(self, z, x, y, z_star, x_star, which_kernels):
        subset = ~torch.isnan(y).reshape(-1)
        if subset.sum() > 0:
            y = y[subset, :]
            z = z[subset, :]
            x = x[subset, :]
        Nstar = z_star.size()[0]
        y = (y - self.intercept)
        sigma2 = self.get_noise_var()

        K_all = self.get_K_without_noise(z, z, x, x, jitter=self.jitter) + sigma2 * torch.eye(y.size()[0])
        L_all = torch.cholesky(K_all)
        K_all_inv = torch.cholesky_inverse(L_all)

        K_sf = self.get_K_without_noise(z_star, z, x_star, x, which_kernels=which_kernels)
        K_ss = self.get_K_without_noise(z_star, z_star, x_star, x_star, which_kernels=which_kernels)

        tmp = torch.mm(K_sf, K_all_inv)
        mean = self.intercept + torch.mm(tmp, y)
        var = torch.diag(K_ss - torch.mm(tmp, K_sf.t()))

        return mean.reshape(-1), var
//  [python]
def predict_decomposition(self, z, x, y, z_star, x_star, which_kernels):
        subset = ~torch.isnan(y).reshape(-1)
        if subset.sum() > 0:
            y = y[subset, :]
            z = z[subset, :]
            x = x[subset, :]
        Nstar = z_star.size()[0]
        y = (y - self.intercept)
        sigma2 = self.get_noise_var()

        K_all = self.get_K_without_noise(z, z, x, x, jitter=self.jitter) + sigma2 * torch.eye(y.size()[0])
        L_all = torch.linalg.cholesky(K_all)
        K_all_inv = torch.linalg.cholesky_inverse(L_all)

        K_sf = self.get_K_without_noise(z_star, z, x_star, x, which_kernels=which_kernels)
        K_ss = self.get_K_without_noise(z_star, z_star, x_star, x_star, which_kernels=which_kernels)

        tmp = torch.mm(K_sf, K_all_inv)
        mean = self.intercept + torch.mm(tmp, y)
        var = torch.diag(K_ss - torch.mm(tmp, K_sf.t()))

        return mean.reshape(-1), var
//  [python]
def _log_likelihood(self, X, k=0, c=None):
        """ Log-probability density function (pdf) of the standard normal
            distribution, evaluated at the values in X.
        Args:
            X (torch.tensor): Observed data (N, C).
            k (int, optional): Index of mixture component. Defaults to 0.
            c (int, optional): Index of channel. Defaults to None.
        Returns:
            log_pdf (torch.tensor): (N, 1).
        """
        C = X.shape[1]
        device = X.device
        dtype = X.dtype
        pi = torch.tensor(math.pi, dtype=dtype, device=device)
        if c is not None:
            Cov = self.Cov[c, c, k].reshape(1, 1).cpu()
            mu = self.mu[c, k].reshape(1).cpu()
        else:
            Cov = self.Cov[:, :, k]
            mu = self.mu[:, k]
        if C == 1:
            chol_Cov = torch.sqrt(Cov)
            log_det_Cov = torch.log(chol_Cov[0, 0])
        else:
            chol_Cov = torch.cholesky(Cov)
            log_det_Cov = torch.sum(torch.log(torch.diag(chol_Cov)))
            chol_Cov = chol_Cov.inverse()
        chol_Cov = chol_Cov.type(dtype)
        mu = mu.type(dtype)
        if C == 1:
            diff = (X - mu)/chol_Cov
        else:
            diff = torch.tensordot(X - mu, chol_Cov, dims=([1], [0]))
        log_pdf = - (C / 2) * torch.log(2 * pi) - log_det_Cov - 0.5 * torch.sum(diff**2, dim=1)
        return log_pdf
//  [python]
def _log_likelihood(self, X, k=0, c=None):
        """ Log-probability density function (pdf) of the standard normal
            distribution, evaluated at the values in X.
        Args:
            X (torch.tensor): Observed data (N, C).
            k (int, optional): Index of mixture component. Defaults to 0.
            c (int, optional): Index of channel. Defaults to None.
        Returns:
            log_pdf (torch.tensor): (N, 1).
        """
        C = X.shape[1]
        device = X.device
        dtype = X.dtype
        pi = torch.tensor(math.pi, dtype=dtype, device=device)
        if c is not None:
            Cov = self.Cov[c, c, k].reshape(1, 1).cpu()
            mu = self.mu[c, k].reshape(1).cpu()
        else:
            Cov = self.Cov[:, :, k]
            mu = self.mu[:, k]
        if C == 1:
            chol_Cov = torch.sqrt(Cov)
            log_det_Cov = torch.log(chol_Cov[0, 0])
        else:
            chol_Cov = torch.linalg.cholesky(Cov)
            log_det_Cov = torch.sum(torch.log(torch.diag(chol_Cov)))
            chol_Cov = chol_Cov.inverse()
        chol_Cov = chol_Cov.type(dtype)
        mu = mu.type(dtype)
        if C == 1:
            diff = (X - mu)/chol_Cov
        else:
            diff = torch.tensordot(X - mu, chol_Cov, dims=([1], [0]))
        log_pdf = - (C / 2) * torch.log(2 * pi) - log_det_Cov - 0.5 * torch.sum(diff**2, dim=1)
        return log_pdf
//  [python]
def log_prob(self, z, x, y):
        subset = ~torch.isnan(y).reshape(-1)
        if subset.sum() > 0:
            y = y[subset, :]
            z = z[subset, :]
            x = x[subset, :]
        # inducing points log_prob
        N = y.size()[0]
        y = (y - self.intercept)
        sigma2 = self.get_noise_var()
        sigma = torch.sqrt(sigma2)
        K_uu = self.get_K_without_noise(self.z_u, self.z_u, self.x_u, self.x_u, jitter=self.jitter)
        K_uf = self.get_K_without_noise(self.z_u, z, self.x_u, x)
        L = torch.cholesky(K_uu)
        A = torch.triangular_solve(K_uf, L, upper=False)[0]
        AAT = torch.mm(A, A.t())
        B = AAT + torch.eye(self.M)
        LB = torch.cholesky(B)
        Aerr = torch.mm(A, y)
        c = torch.triangular_solve(Aerr, LB, upper=False)[0] / sigma

        Kdiag = self.get_K_diag(z, x)

        bound = -0.5 * N * np.log(2 * np.pi) - torch.sum(torch.log(torch.diag(LB))) - 0.5 * N * torch.log(
            sigma2) - 0.5 * torch.sum(torch.pow(y, 2)) / sigma2
        bound += 0.5 * torch.sum(torch.pow(c, 2)) - 0.5 * torch.sum(Kdiag) / sigma2 + 0.5 * torch.sum(torch.diag(AAT))
        return bound