jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def _apply_non_linear_constraint_at_t(  # noqa: WPS210, WPS231
    loc_t: Tensor, cov_t: Tensor, max_iter: int, tol: float, mode: InequalityConstraint
):
    nx = loc_t.shape[0]
    eye = torch.eye(nx).to(loc_t.device)
    chol_t: Tensor = torch.linalg.cholesky(cov_t)
    chinv: Tensor = torch.linalg.solve_triangular(chol_t, eye, upper=False)
    chinv_loc_t = chinv.matmul(loc_t)
    if mode == "nonneg":
        a_matrix = chinv.numpy()
        b = chinv_loc_t.squeeze(-1).numpy()
        loc_t_np, _ = nnls(a_matrix, b)
        return torch.from_numpy(loc_t_np).unsqueeze(-1)
    if mode == "monotone":
        diff_matrix = _construct_diff_matrix(nx)
        diff_matrix_inv = np.linalg.inv(diff_matrix)
        a_matrix = np.matmul(chinv.numpy(), diff_matrix_inv)
        b = chinv_loc_t.squeeze(-1).numpy()
        loc_t_np_diff, _ = nnls(a_matrix, b)
        loc_t_np = np.matmul(diff_matrix_inv, loc_t_np_diff)
        return torch.from_numpy(loc_t_np).unsqueeze(-1)
//  [python]
def orth(X):
    S = X @ X.mT
    eps = S.diagonal(dim1=1, dim2=2).mean(1).mul(1e-3).detach()
    eye = torch.eye(S.size(-1), dtype=S.dtype, device=S.device)
    S = S + eps.view(-1, 1, 1) * eye.unsqueeze(0)
    L = torch.cholesky(S)
    W = torch.linalg.solve_triangular(L, X, upper=False)
    return W
//  [python]
def orth(X):
    S = X @ X.mT
    eps = S.diagonal(dim1=1, dim2=2).mean(1).mul(1e-3).detach()
    eye = torch.eye(S.size(-1), dtype=S.dtype, device=S.device)
    S = S + eps.view(-1, 1, 1) * eye.unsqueeze(0)
    L = torch.linalg.cholesky(S)
    W = torch.linalg.solve_triangular(L, X, upper=False)
    return W
//  [python]
def forward(ctx, X):
        S = X @ X.mT
        eps = S.diagonal(dim1=1, dim2=2).mean(1).mul(1e-3)
        eye = torch.eye(S.size(-1), dtype=S.dtype, device=S.device)
        S = S + eps.view(-1, 1, 1) * eye.unsqueeze(0)
        L = torch.cholesky(S)
        W = torch.linalg.solve_triangular(L, X, upper=False)
        ctx.save_for_backward(W, L)
        return W
//  [python]
def forward(ctx, X):
        S = X @ X.mT
        eps = S.diagonal(dim1=1, dim2=2).mean(1).mul(1e-3)
        eye = torch.eye(S.size(-1), dtype=S.dtype, device=S.device)
        S = S + eps.view(-1, 1, 1) * eye.unsqueeze(0)
        L = torch.linalg.cholesky(S)
        W = torch.linalg.solve_triangular(L, X, upper=False)
        ctx.save_for_backward(W, L)
        return W
//  [python]
def get_prior_matrices(self):
        """"
        """
        Ku = self.kern.K(self.inducing_points)
        Lu = Ku.cholesky()
        Ku_tiled = torch.tile(Ku[None, :, :], [self.num_outputs, 1, 1])
        Lu_tiled = torch.tile(Lu[None, :, :], [self.num_outputs, 1, 1])


        return Ku, Lu, Ku_tiled, Lu_tiled
//  [python]
def get_prior_matrices(self):
        """"
        """
        Ku = self.kern.K(self.inducing_points)
        Lu = torch.linalg.cholesky(Ku)# Ku.cholesky()
        Ku_tiled = torch.tile(Ku[None, :, :], [self.num_outputs, 1, 1])
        Lu_tiled = torch.tile(Lu[None, :, :], [self.num_outputs, 1, 1])


        return Ku, Lu, Ku_tiled, Lu_tiled
//  [python]
def psd_safe_cholesky(A, upper=False, out=None, jitter=None):
	"""Compute the Cholesky decomposition of A. If A is only p.s.d, add a small jitter to the diagonal.
	Args:
		:attr:`A` (Tensor):
			The tensor to compute the Cholesky decomposition of
		:attr:`upper` (bool, optional):
			See torch.cholesky
		:attr:`out` (Tensor, optional):
			See torch.cholesky
		:attr:`jitter` (float, optional):
			The jitter to add to the diagonal of A in case A is only p.s.d. If omitted, chosen
			as 1e-6 (float) or 1e-8 (double)
	"""
	try:
		L = torch.cholesky(A, upper=upper, out=out)
		return L
	except RuntimeError as e:
		isnan = torch.isnan(A)
		if isnan.any():
			raise NanError(
				f"cholesky_cpu: {isnan.sum().item()} of {A.numel()} elements of the {A.shape} tensor are NaN."
			)

		if jitter is None:
			jitter = 1e-6 if A.dtype == torch.float32 else 1e-8
		Aprime = A.clone()
		jitter_prev = 0
		for i in range(10):
			jitter_new = jitter * (10 ** i)
			Aprime.diagonal(dim1=-2, dim2=-1).add_(jitter_new - jitter_prev)
			jitter_prev = jitter_new
			try:
				L = torch.cholesky(Aprime, upper=upper, out=out)
				warnings.warn(
					f"A not p.d., added jitter of {jitter_new} to the diagonal",
					RuntimeWarning,
				)
				return L
			except RuntimeError:
				continue
		# return torch.randn_like(A).tril()
		raise e
//  [python]
def psd_safe_cholesky(A, upper=False, out=None, jitter=None):
	"""Compute the Cholesky decomposition of A. If A is only p.s.d, add a small jitter to the diagonal.
	Args:
		:attr:`A` (Tensor):
			The tensor to compute the Cholesky decomposition of
		:attr:`upper` (bool, optional):
			See torch.cholesky
		:attr:`out` (Tensor, optional):
			See torch.cholesky
		:attr:`jitter` (float, optional):
			The jitter to add to the diagonal of A in case A is only p.s.d. If omitted, chosen
			as 1e-6 (float) or 1e-8 (double)
	"""
	try:
		L = torch.linalg.cholesky(A, upper=upper, out=out)
		return L
	except RuntimeError as e:
		isnan = torch.isnan(A)
		if isnan.any():
			raise NanError(
				f"cholesky_cpu: {isnan.sum().item()} of {A.numel()} elements of the {A.shape} tensor are NaN."
			)

		if jitter is None:
			jitter = 1e-6 if A.dtype == torch.float32 else 1e-8
		Aprime = A.clone()
		jitter_prev = 0
		for i in range(10):
			jitter_new = jitter * (10 ** i)
			Aprime.diagonal(dim1=-2, dim2=-1).add_(jitter_new - jitter_prev)
			jitter_prev = jitter_new
			try:
				L = torch.linalg.cholesky(Aprime, upper=upper, out=out)
				warnings.warn(
					f"A not p.d., added jitter of {jitter_new} to the diagonal",
					RuntimeWarning,
				)
				return L
			except RuntimeError:
				continue
		# return torch.randn_like(A).tril()
		raise e
//  [python]
def test_cholesky_inverse(self, device, dtype):
        from torch.testing._internal.common_utils import random_hermitian_pd_matrix

        def run_test(shape, batch, upper, contiguous):
            A = random_hermitian_pd_matrix(shape, *batch, dtype=dtype, device=device)
            if A.numel() > 0 and not contiguous:
                A = A.transpose(-2, -1)
                self.assertFalse(A.is_contiguous())
            L = torch.cholesky(A)
            expected_inverse = torch.inverse(A)
            L = L.conj().transpose(-2, -1) if upper else L
            actual_inverse = torch.cholesky_inverse(L, upper)
            self.assertEqual(actual_inverse, expected_inverse)

        shapes = (0, 3, 5)
        batches = ((), (0,), (3, ), (2, 2))
        for shape, batch, upper, contiguous in list(itertools.product(shapes, batches, (True, False), (True, False))):
            run_test(shape, batch, upper, contiguous)

        # check the out= variant
        A = random_hermitian_pd_matrix(3, 2, dtype=dtype, device=device)
        L = torch.cholesky(A)

        # There are two code paths currently for the out= variant
        # 1. When 'out' tensor is in Fortran (column-major) memory format
        # then the fast route is taken and the storage is reused directly in the computations
        # 2. When 'out' tensor is not in Fortran format then a temporary tensor is allocated internally
        # and the result is copied from the temporary tensor to 'out' tensor

        # This test checks the first code path
        out = torch.empty_like(A)
        out_t = out.transpose(-2, -1).clone(memory_format=torch.contiguous_format)
        out = out_t.transpose(-2, -1)
        ans = torch.cholesky_inverse(L, out=out)
        self.assertEqual(ans, out)
        expected = torch.inverse(A)
        self.assertEqual(expected, out)

        # This test checks the second code path
        out = torch.empty_like(A)
        ans = torch.cholesky_inverse(L, out=out)
        self.assertEqual(ans, out)
        expected = torch.inverse(A)
        self.assertEqual(expected, out)
//  [python]
def test_cholesky_inverse(self, device, dtype):
        from torch.testing._internal.common_utils import random_hermitian_pd_matrix

        def run_test(shape, batch, upper, contiguous):
            A = random_hermitian_pd_matrix(shape, *batch, dtype=dtype, device=device)
            if A.numel() > 0 and not contiguous:
                A = A.transpose(-2, -1)
                self.assertFalse(A.is_contiguous())
            L = torch.linalg.cholesky(A)
            expected_inverse = torch.inverse(A)
            L = L.conj().transpose(-2, -1) if upper else L
            actual_inverse = torch.cholesky_inverse(L, upper)
            self.assertEqual(actual_inverse, expected_inverse)

        shapes = (0, 3, 5)
        batches = ((), (0,), (3, ), (2, 2))
        for shape, batch, upper, contiguous in list(itertools.product(shapes, batches, (True, False), (True, False))):
            run_test(shape, batch, upper, contiguous)

        # check the out= variant
        A = random_hermitian_pd_matrix(3, 2, dtype=dtype, device=device)
        L = torch.linalg.cholesky(A)

        # There are two code paths currently for the out= variant
        # 1. When 'out' tensor is in Fortran (column-major) memory format
        # then the fast route is taken and the storage is reused directly in the computations
        # 2. When 'out' tensor is not in Fortran format then a temporary tensor is allocated internally
        # and the result is copied from the temporary tensor to 'out' tensor

        # This test checks the first code path
        out = torch.empty_like(A)
        out_t = out.transpose(-2, -1).clone(memory_format=torch.contiguous_format)
        out = out_t.transpose(-2, -1)
        ans = torch.cholesky_inverse(L, out=out)
        self.assertEqual(ans, out)
        expected = torch.inverse(A)
        self.assertEqual(expected, out)

        # This test checks the second code path
        out = torch.empty_like(A)
        ans = torch.cholesky_inverse(L, out=out)
        self.assertEqual(ans, out)
        expected = torch.inverse(A)
        self.assertEqual(expected, out)
//  [python]
def em_update(self, _triggered=False):
        """
        Perform a modified em-update for MVG, updating the cholesky decomposition (CD) of the covariance matrix using
        the CD from stats_acc.
        For the remainder, the standard update method is called.

        :param _triggered: for internal use, don't set
        :return: None
        """

        if self._online_em_stepsize is not None and not _triggered:
            return

        mu_update = self._stats_acc[..., :self.num_dims].clone()
        cov_update = self._stats_acc[..., self.num_dims:].clone().reshape((*self._stats_acc.size()[:-1],
                                                                           self.num_dims, self.num_dims))

        # Add a small value to the diagonal to ensure positive eigenvalues
        cov_update[..., range(self.num_dims), range(self.num_dims)] += 1e-3

        # Calculate CD and flatten it
        l = torch.cholesky(cov_update)
        tri_indices = torch.tril_indices(self.num_dims, self.num_dims, offset=-1, device=l.device)
        l = torch.cat([l.diagonal(dim1=-2, dim2=-1), l[..., tri_indices[0], tri_indices[1]]], dim=-1)

        stats_acc = torch.zeros(
            *self._stats_acc.size()[:-1], self.num_dims + self.cov_cd_size).to(self._stats_acc.device)
        stats_acc[..., :self.num_dims] = mu_update
        stats_acc[..., self.num_dims:] = l

        self._stats_acc = stats_acc

        super().em_update(_triggered)
//  [python]
def em_update(self, _triggered=False):
        """
        Perform a modified em-update for MVG, updating the cholesky decomposition (CD) of the covariance matrix using
        the CD from stats_acc.
        For the remainder, the standard update method is called.

        :param _triggered: for internal use, don't set
        :return: None
        """

        if self._online_em_stepsize is not None and not _triggered:
            return

        mu_update = self._stats_acc[..., :self.num_dims].clone()
        cov_update = self._stats_acc[..., self.num_dims:].clone().reshape((*self._stats_acc.size()[:-1],
                                                                           self.num_dims, self.num_dims))

        # Add a small value to the diagonal to ensure positive eigenvalues
        cov_update[..., range(self.num_dims), range(self.num_dims)] += 1e-3

        # Calculate CD and flatten it
        l = torch.linalg.cholesky(cov_update)
        tri_indices = torch.tril_indices(self.num_dims, self.num_dims, offset=-1, device=l.device)
        l = torch.cat([l.diagonal(dim1=-2, dim2=-1), l[..., tri_indices[0], tri_indices[1]]], dim=-1)

        stats_acc = torch.zeros(
            *self._stats_acc.size()[:-1], self.num_dims + self.cov_cd_size).to(self._stats_acc.device)
        stats_acc[..., :self.num_dims] = mu_update
        stats_acc[..., self.num_dims:] = l

        self._stats_acc = stats_acc

        super().em_update(_triggered)
//  [python]
def match_histogram(target, source, mode="chol", eps=1e-2):
    """from https://github.com/ProGamerGov/Neural-Tools/blob/master/linear-color-transfer.py#L36"""

    target = target.permute(0, 3, 1, 2)  # -> b, c, h, w
    source = source.permute(0, 3, 1, 2)

    mu_t = target.mean((2, 3), keepdim=True)
    hist_t = (target - mu_t).view(target.size(1), -1)  # [c, b * h * w]
    cov_t = hist_t @ hist_t.T / hist_t.shape[1] + eps * torch.eye(hist_t.shape[0], device=device)

    mu_s = source.mean((2, 3), keepdim=True)
    hist_s = (source - mu_s).view(source.size(1), -1)
    cov_s = hist_s @ hist_s.T / hist_s.shape[1] + eps * torch.eye(hist_s.shape[0], device=device)

    if mode == "chol":
        chol_t = torch.cholesky(cov_t)
        chol_s = torch.cholesky(cov_s)
        matched = chol_s @ torch.inverse(chol_t) @ hist_t

    elif mode == "pca":
        eva_t, eve_t = torch.linalg.eigh(cov_t, UPLO="U")
        Qt = eve_t @ torch.sqrt(torch.diag(eva_t)) @ eve_t.T
        eva_s, eve_s = torch.linalg.eigh(cov_s, UPLO="U")
        Qs = eve_s @ torch.sqrt(torch.diag(eva_s)) @ eve_s.T
        matched = Qs @ torch.inverse(Qt) @ hist_t

    elif mode == "sym":
        eva_t, eve_t = torch.linalg.eigh(cov_t, UPLO="U")
        Qt = eve_t @ torch.sqrt(torch.diag(eva_t)) @ eve_t.T
        Qt_Cs_Qt = Qt @ cov_s @ Qt
        eva_QtCsQt, eve_QtCsQt = torch.linalg.eigh(Qt_Cs_Qt, UPLO="U")
        QtCsQt = eve_QtCsQt @ torch.sqrt(torch.diag(eva_QtCsQt)) @ eve_QtCsQt.T
        matched = torch.inverse(Qt) @ QtCsQt @ torch.inverse(Qt) @ hist_t

    matched = matched.view(*target.shape) + mu_s

    return matched.permute(0, 2, 3, 1)
//  [python]
def match_histogram(target, source, mode="chol", eps=1e-2):
    """from https://github.com/ProGamerGov/Neural-Tools/blob/master/linear-color-transfer.py#L36"""

    target = target.permute(0, 3, 1, 2)  # -> b, c, h, w
    source = source.permute(0, 3, 1, 2)

    mu_t = target.mean((2, 3), keepdim=True)
    hist_t = (target - mu_t).view(target.size(1), -1)  # [c, b * h * w]
    cov_t = hist_t @ hist_t.T / hist_t.shape[1] + eps * torch.eye(hist_t.shape[0], device=device)

    mu_s = source.mean((2, 3), keepdim=True)
    hist_s = (source - mu_s).view(source.size(1), -1)
    cov_s = hist_s @ hist_s.T / hist_s.shape[1] + eps * torch.eye(hist_s.shape[0], device=device)

    if mode == "chol":
        chol_t = torch.linalg.cholesky(cov_t)
        chol_s = torch.linalg.cholesky(cov_s)
        matched = chol_s @ torch.inverse(chol_t) @ hist_t

    elif mode == "pca":
        eva_t, eve_t = torch.linalg.eigh(cov_t, UPLO="U")
        Qt = eve_t @ torch.sqrt(torch.diag(eva_t)) @ eve_t.T
        eva_s, eve_s = torch.linalg.eigh(cov_s, UPLO="U")
        Qs = eve_s @ torch.sqrt(torch.diag(eva_s)) @ eve_s.T
        matched = Qs @ torch.inverse(Qt) @ hist_t

    elif mode == "sym":
        eva_t, eve_t = torch.linalg.eigh(cov_t, UPLO="U")
        Qt = eve_t @ torch.sqrt(torch.diag(eva_t)) @ eve_t.T
        Qt_Cs_Qt = Qt @ cov_s @ Qt
        eva_QtCsQt, eve_QtCsQt = torch.linalg.eigh(Qt_Cs_Qt, UPLO="U")
        QtCsQt = eve_QtCsQt @ torch.sqrt(torch.diag(eva_QtCsQt)) @ eve_QtCsQt.T
        matched = torch.inverse(Qt) @ QtCsQt @ torch.inverse(Qt) @ hist_t

    matched = matched.view(*target.shape) + mu_s

    return matched.permute(0, 2, 3, 1)
//  [python]
def __init__(self, kern, Z, num_outputs, mean_function,
                 white=False, input_prop_dim=None, seed=0, dtype = torch.float64, device = None, ind_total_nan=None, n_layers=None):
        """
        """
        super().__init__(input_prop_dim=input_prop_dim, dtype = dtype, device=device)
        self.seed = seed

        self.num_inducing = Z.shape[0]

        # Inducing points prior mean
        self.register_parameter('q_mu', nn.Parameter(
            torch.tensor(np.zeros((self.num_inducing, num_outputs)), dtype=self.dtype, device=self.device)))


        self.register_parameter('inducing_points', nn.Parameter(Z.clone().detach().requires_grad_(True)))


        self.kern = kern
        self.mean_function = mean_function

        self.num_outputs = num_outputs

        self.white = white
        self.ind_total_nan = ind_total_nan


        if not self.white:  # initialize to prior
            Ku = self.kern.K(Z)
            Lu = torch.cholesky(Ku+torch.eye(Z.shape[0]).to(self.device)*default_jitter)
            Lu_tiled = torch.tile(Lu[None, :, :], [num_outputs, 1, 1])
            li, lj = torch.tril_indices(self.num_inducing, self.num_inducing, 0)
            triangular_Lu = Lu_tiled[:, li, lj]

            self.register_parameter('q_sqrt_tri', nn.Parameter(triangular_Lu.clone().detach().requires_grad_(True)))
//  [python]
def __init__(self, kern, Z, num_outputs, mean_function,
                 white=False, input_prop_dim=None, seed=0, dtype = torch.float64, device = None, ind_total_nan=None, n_layers=None):
        """
        """
        super().__init__(input_prop_dim=input_prop_dim, dtype = dtype, device=device)
        self.seed = seed

        self.num_inducing = Z.shape[0]

        # Inducing points prior mean
        self.register_parameter('q_mu', nn.Parameter(
            torch.tensor(np.zeros((self.num_inducing, num_outputs)), dtype=self.dtype, device=self.device)))


        self.register_parameter('inducing_points', nn.Parameter(Z.clone().detach().requires_grad_(True)))


        self.kern = kern
        self.mean_function = mean_function

        self.num_outputs = num_outputs

        self.white = white
        self.ind_total_nan = ind_total_nan


        if not self.white:  # initialize to prior
            Ku = self.kern.K(Z)
            Lu = torch.linalg.cholesky(Ku+torch.eye(Z.shape[0]).to(self.device)*default_jitter)
            Lu_tiled = torch.tile(Lu[None, :, :], [num_outputs, 1, 1])
            li, lj = torch.tril_indices(self.num_inducing, self.num_inducing, 0)
            triangular_Lu = Lu_tiled[:, li, lj]

            self.register_parameter('q_sqrt_tri', nn.Parameter(triangular_Lu.clone().detach().requires_grad_(True)))
//  [python]
def wasserstein_distance_mvn(mean1, cov1, mean2, cov2):
	# calculate Tr((cov1 * cov2)^(1/2))
	sq_tr_cov = torch.trace(torch.cholesky(torch.matmul(cov1, cov2)))
	# plug the sqrt_trace_component into Tr(cov1 + cov2 - 2(cov1 * cov2)^(1/2))
	trace_term = torch.trace(cov1 + cov2) - 2.0 * sq_tr_cov  # scalar
	# |mean1 - mean2|^2
	diff = mean1 - mean2  # [n, 1]
	mean_term = torch.sum(torch.mul(diff, diff))  # scalar
	# put it together
	return (trace_term + mean_term).float()
//  [python]
def wasserstein_distance_mvn(mean1, cov1, mean2, cov2):
	# calculate Tr((cov1 * cov2)^(1/2))
	sq_tr_cov = torch.trace(torch.linalg.cholesky(torch.matmul(cov1, cov2)))
	# plug the sqrt_trace_component into Tr(cov1 + cov2 - 2(cov1 * cov2)^(1/2))
	trace_term = torch.trace(cov1 + cov2) - 2.0 * sq_tr_cov  # scalar
	# |mean1 - mean2|^2
	diff = mean1 - mean2  # [n, 1]
	mean_term = torch.sum(torch.mul(diff, diff))  # scalar
	# put it together
	return (trace_term + mean_term).float()
//  [python]
def _self_inverse(self, matrix):
        temp_matrix = torch.linalg.inv(torch.cholesky(matrix))
        return temp_matrix @ temp_matrix.T
//  [python]
def _self_inverse(self, matrix):
        temp_matrix = torch.linalg.inv(torch.linalg.cholesky(matrix))
        return temp_matrix @ temp_matrix.T
//  [python]
def draw_perfect(self,M):
        global ep, dp
        cols = self.generateRVIndices()
        # draw it from prior first
        theta = np.zeros((M,self.dim()))
        # hack it
        theta[:] = self.sampleFromPrior(M)
        print(theta)
        # That sorts out model indicators for now.
        # FIXME for now we leave model indicators as drawn from prior
        # and just draw perfect draws for these models again.
        # This doesn't give us a complete joint posterior
        # because the model probability marginals are wrong, they're the prior.
        # For calibration, we don't care.

        k = theta[:,cols['k']].flatten()
        t1 = theta[:,cols['t1']].flatten()
        t2 = theta[:,cols['t2']].flatten()
        m1 = k==0
        m2 = k==1
        mk1cov = 1.
        mk2cov = torch.tensor([[1.,.99],[.99,1.]]) 
        if True:
            tf1  = n.transforms.InverseTransform(SinArcSinhTransform(ep[0],dp[0]))
            tf2  = n.transforms.InverseTransform(n.transforms.CompositeTransform([LTransform(torch.cholesky(torch.tensor([[1.,.99],[.99,1.]]))),SAS2DTransform([ep[1],ep[2]],[dp[1],dp[2]])]))
            #tf1  = n.transforms.InverseTransform(L1DTransform(np.sqrt(mk1cov)))
            #tf1  = n.transforms.IdentityTransform()
            #tf2  = n.transforms.InverseTransform(LTransform(torch.cholesky(mk2cov)))
            f1 = Flow(tf1,StandardNormal((1,)))
            f2 = Flow(tf2,StandardNormal((2,)))
            m1_v = f1.sample(t1[m1].shape[0]).detach().numpy()
            t1[m1]= m1_v[:,0]
            m2_v = f2.sample(t1[m2].shape[0]).detach().numpy()
            t1[m2]=m2_v[:,0]
            t2[m2]=m2_v[:,1]
            theta[:,cols['t1']] = t1.reshape(theta[:,cols['t1']].shape)
            theta[:,cols['t2']] = t2.reshape(theta[:,cols['t2']].shape)
        return theta
//  [python]
def draw_perfect(self,M):
        global ep, dp
        cols = self.generateRVIndices()
        # draw it from prior first
        theta = np.zeros((M,self.dim()))
        # hack it
        theta[:] = self.sampleFromPrior(M)
        print(theta)
        # That sorts out model indicators for now.
        # FIXME for now we leave model indicators as drawn from prior
        # and just draw perfect draws for these models again.
        # This doesn't give us a complete joint posterior
        # because the model probability marginals are wrong, they're the prior.
        # For calibration, we don't care.

        k = theta[:,cols['k']].flatten()
        t1 = theta[:,cols['t1']].flatten()
        t2 = theta[:,cols['t2']].flatten()
        m1 = k==0
        m2 = k==1
        mk1cov = 1.
        mk2cov = torch.tensor([[1.,.99],[.99,1.]]) 
        if True:
            tf1  = n.transforms.InverseTransform(SinArcSinhTransform(ep[0],dp[0]))
            tf2  = n.transforms.InverseTransform(n.transforms.CompositeTransform([LTransform(torch.linalg.cholesky(torch.tensor([[1.,.99],[.99,1.]]))),SAS2DTransform([ep[1],ep[2]],[dp[1],dp[2]])]))
            #tf1  = n.transforms.InverseTransform(L1DTransform(np.sqrt(mk1cov)))
            #tf1  = n.transforms.IdentityTransform()
            #tf2  = n.transforms.InverseTransform(LTransform(torch.linalg.cholesky(mk2cov)))
            f1 = Flow(tf1,StandardNormal((1,)))
            f2 = Flow(tf2,StandardNormal((2,)))
            m1_v = f1.sample(t1[m1].shape[0]).detach().numpy()
            t1[m1]= m1_v[:,0]
            m2_v = f2.sample(t1[m2].shape[0]).detach().numpy()
            t1[m2]=m2_v[:,0]
            t2[m2]=m2_v[:,1]
            theta[:,cols['t1']] = t1.reshape(theta[:,cols['t1']].shape)
            theta[:,cols['t2']] = t2.reshape(theta[:,cols['t2']].shape)
        return theta
//  [python]
def test_update_(self):
        # Construct permuted matrices A
        n = self.matrix.shape[-1]
        A = (1 / self.diag).unsqueeze(-1) * self.matrix * (1 / self.diag).unsqueeze(-2)
        A = A.gather(-1, self.piv_chol.perm.unsqueeze(-2).repeat(1, n, 1))
        A = A.gather(-2, self.piv_chol.perm.unsqueeze(-1).repeat(1, 1, n))

        # Test upper left block
        L = torch.cholesky(A[..., :2, :2])
        self.assertTrue(L.allclose(self.piv_chol.tril[..., :2, :2]))

        # Test lower left block
        beta = torch.triangular_solve(A[..., :2:, 2:], L, upper=False).solution
        self.assertTrue(
            beta.transpose(-1, -2).allclose(self.piv_chol.tril[..., 2:, :2])
        )

        # Test lower right block
        schur = A[..., 2:, 2:] - beta.transpose(-1, -2) @ beta
        self.assertTrue(schur.tril().allclose(self.piv_chol.tril[..., 2:, 2:]))
//  [python]
def test_update_(self):
        # Construct permuted matrices A
        n = self.matrix.shape[-1]
        A = (1 / self.diag).unsqueeze(-1) * self.matrix * (1 / self.diag).unsqueeze(-2)
        A = A.gather(-1, self.piv_chol.perm.unsqueeze(-2).repeat(1, n, 1))
        A = A.gather(-2, self.piv_chol.perm.unsqueeze(-1).repeat(1, 1, n))

        # Test upper left block
        L = torch.linalg.cholesky(A[..., :2, :2])
        self.assertTrue(L.allclose(self.piv_chol.tril[..., :2, :2]))

        # Test lower left block
        beta = torch.linalg.solve_triangular(L, A[..., :2:, 2:], upper=False)
        self.assertTrue(
            beta.transpose(-1, -2).allclose(self.piv_chol.tril[..., 2:, :2])
        )

        # Test lower right block
        schur = A[..., 2:, 2:] - beta.transpose(-1, -2) @ beta
        self.assertTrue(schur.tril().allclose(self.piv_chol.tril[..., 2:, 2:]))
//  [python]
def test_augment(self):
        K = self.matrix
        n = K.shape[-1]
        m = n // 2
        Kaa = K[:, 0:m, 0:m]
        Laa = t.cholesky(Kaa)
        Kbb = K[:, m:, m:]

        error_msg = "One and only one of `Kba` or `Lba` must be provided."
        with self.assertRaisesRegex(ValueError, error_msg):
            augment_cholesky(Laa, Kbb)

        Kba = K[:, m:, 0:m]
        L_augmented = augment_cholesky(Laa, Kbb, Kba)
        L = t.cholesky(K)
        self.assertAllClose(L_augmented, L)

        # with jitter
        jitter = 3e-2
        Laa = t.cholesky(Kaa + jitter * torch.eye(m).unsqueeze(0))
        L_augmented = augment_cholesky(Laa, Kbb, Kba, jitter=jitter)
        L = t.cholesky(K + jitter * torch.eye(n).unsqueeze(0))
        self.assertAllClose(L_augmented, L)
//  [python]
def test_augment(self):
        K = self.matrix
        n = K.shape[-1]
        m = n // 2
        Kaa = K[:, 0:m, 0:m]
        Laa = torch.linalg.cholesky(Kaa)
        Kbb = K[:, m:, m:]

        error_msg = "One and only one of `Kba` or `Lba` must be provided."
        with self.assertRaisesRegex(ValueError, error_msg):
            augment_cholesky(Laa, Kbb)

        Kba = K[:, m:, 0:m]
        L_augmented = augment_cholesky(Laa, Kbb, Kba)
        L = torch.linalg.cholesky(K)
        self.assertAllClose(L_augmented, L)

        # with jitter
        jitter = 3e-2
        Laa = torch.linalg.cholesky(Kaa + jitter * torch.eye(m).unsqueeze(0))
        L_augmented = augment_cholesky(Laa, Kbb, Kba, jitter=jitter)
        L = torch.linalg.cholesky(K + jitter * torch.eye(n).unsqueeze(0))
        self.assertAllClose(L_augmented, L)
//  [python]
def negative_log_likelihood(self):
        y_num, y_dimension = self.Y.shape
        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))

        L = torch.cholesky(Sigma)
        #option 1 (use this if torch supports)
        Gamma,_ = torch.triangular_solve(self.Y - self.meanFunc(self.X), L, upper = False)
        #option 2
        # gamma = L.inverse() @ Y       # we can use this as an alternative because L is a lower triangular matrix.

        nll =  0.5 * (Gamma ** 2).sum() +  L.diag().log().sum() * y_dimension  \
            + 0.5 * y_num * torch.log(2 * torch.tensor(PI)) * y_dimension
        return nll
//  [python]
def negative_log_likelihood(self):
        y_num, y_dimension = self.Y.shape
        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))

        L = torch.linalg.cholesky(Sigma)
        #option 1 (use this if torch supports)
        Gamma,_ = torch.triangular_solve(self.Y - self.meanFunc(self.X), L, upper = False)
        #option 2
        # gamma = L.inverse() @ Y       # we can use this as an alternative because L is a lower triangular matrix.

        nll =  0.5 * (Gamma ** 2).sum() +  L.diag().log().sum() * y_dimension  \
            + 0.5 * y_num * torch.log(2 * torch.tensor(PI)) * y_dimension
        return nll
//  [python]
def forward(self, x1):
        v1 = self.linear(x1)
        v2 = v1 - 0.75
        v3 = v1 - 0.2857142857142857
        v4 = v1 + 0.0357142857142857
        v5 = -v1 - 0.7142857142857143
        v6 = v2 * v3 
        v7 = v2 + v3
        v8 = torch.cosine_similarity(v4, v5)
        v9 = torch.addcmul(v6, v4, v5)
        v10 = torch.addmm(v7, v4, v5)
        v11 = torch.cholesky(v8)
        return v9 + v10 + v11
//  [python]
def forward(self, x1):
        v1 = self.linear(x1)
        v2 = v1 - 0.75
        v3 = v1 - 0.2857142857142857
        v4 = v1 + 0.0357142857142857
        v5 = -v1 - 0.7142857142857143
        v6 = v2 * v3 
        v7 = v2 + v3
        v8 = torch.cosine_similarity(v4, v5)
        v9 = torch.addcmul(v6, v4, v5)
        v10 = torch.addmm(v7, v4, v5)
        v11 = torch.linalg.cholesky(v8)
        return v9 + v10 + v11
//  [python]
def compute_llh(self,theta):
        global ep, dp, m1prob
        # we use the RJ target from Andrieu et al 2009
        # \pi(\theta,k) = 0.25 * N(\theta;0,1) * I(k==1) + 0.75 * N(\theta;[0,0],[[1,-0.9],[-0.9,1]]) * I(k==2)
        # k \in {1,2}
        cols = self.generateRVIndices()
        k = theta[:,cols['k']].flatten()
        t1 = theta[:,cols['t1']].flatten()
        t2 = theta[:,cols['t2']].flatten()
        m1 = k==0
        m2 = k==1
        llh = np.zeros(k.shape[0])
        mk1cov = 1.
        mk2cov = torch.tensor([[1.,.99],[.99,1.]]) 
        if True:
            tf1  = n.transforms.InverseTransform(SinArcSinhTransform(ep[0],dp[0]))
            tf2  = n.transforms.InverseTransform(n.transforms.CompositeTransform([LTransform(torch.cholesky(torch.tensor([[1.,.99],[.99,1.]]))),SAS2DTransform([ep[1],ep[2]],[dp[1],dp[2]])]))
            f1 = Flow(tf1,StandardNormal((1,)))
            f2 = Flow(tf2,StandardNormal((2,)))
            if m1.sum() > 0:
                llh[m1] = np.log(m1prob) + f1.log_prob(torch.Tensor(t1[m1].reshape((t1[m1].shape[0],1)))).detach().numpy().flatten()
            if m2.sum() > 0:
                llh[m2] = np.log(1-m1prob) + f2.log_prob(torch.Tensor(np.column_stack([t1[m2],t2[m2]]))).detach().numpy().flatten()
        return llh
//  [python]
def compute_llh(self,theta):
        global ep, dp, m1prob
        # we use the RJ target from Andrieu et al 2009
        # \pi(\theta,k) = 0.25 * N(\theta;0,1) * I(k==1) + 0.75 * N(\theta;[0,0],[[1,-0.9],[-0.9,1]]) * I(k==2)
        # k \in {1,2}
        cols = self.generateRVIndices()
        k = theta[:,cols['k']].flatten()
        t1 = theta[:,cols['t1']].flatten()
        t2 = theta[:,cols['t2']].flatten()
        m1 = k==0
        m2 = k==1
        llh = np.zeros(k.shape[0])
        mk1cov = 1.
        mk2cov = torch.tensor([[1.,.99],[.99,1.]]) 
        if True:
            tf1  = n.transforms.InverseTransform(SinArcSinhTransform(ep[0],dp[0]))
            tf2  = n.transforms.InverseTransform(n.transforms.CompositeTransform([LTransform(torch.linalg.cholesky(torch.tensor([[1.,.99],[.99,1.]]))),SAS2DTransform([ep[1],ep[2]],[dp[1],dp[2]])]))
            f1 = Flow(tf1,StandardNormal((1,)))
            f2 = Flow(tf2,StandardNormal((2,)))
            if m1.sum() > 0:
                llh[m1] = np.log(m1prob) + f1.log_prob(torch.Tensor(t1[m1].reshape((t1[m1].shape[0],1)))).detach().numpy().flatten()
            if m2.sum() > 0:
                llh[m2] = np.log(1-m1prob) + f2.log_prob(torch.Tensor(np.column_stack([t1[m2],t2[m2]]))).detach().numpy().flatten()
        return llh
//  [python]
def draw_sample(self, mu, cov, eps=1e-10):
        L = torch.cholesky(cov + eps * torch.eye(cov.shape[0], device=device))
        return mu + L @ torch.randn(L.shape[0], device=device)
//  [python]
def draw_sample(self, mu, cov, eps=1e-10):
        L = torch.linalg.cholesky(cov + eps * torch.eye(cov.shape[0], device=device))
        return mu + L @ torch.randn(L.shape[0], device=device)
//  [python]
def test_precision_to_scale_tril(batch_shape, event_shape):
    x = torch.randn(batch_shape + event_shape + event_shape)
    precision = x.matmul(x.transpose(-2, -1))
    actual = precision_to_scale_tril(precision)
    expected = torch.cholesky(precision.inverse())
    assert_close(actual, expected)
//  [python]
def test_precision_to_scale_tril(batch_shape, event_shape):
    x = torch.randn(batch_shape + event_shape + event_shape)
    precision = x.matmul(x.transpose(-2, -1))
    actual = precision_to_scale_tril(precision)
    expected = torch.linalg.cholesky(precision.inverse())
    assert_close(actual, expected)
//  [python]
def compute_llh(self,theta):
        global ep, dp, m1prob
        # we use the RJ target from Andrieu et al 2009
        # \pi(\theta,k) = 0.25 * N(\theta;0,1) * I(k==1) + 0.75 * N(\theta;[0,0],[[1,-0.9],[-0.9,1]]) * I(k==2)
        # k \in {1,2}
        # What kind of priors do I use here? Seems they are ill-formed. Perhaps really wide uniform priors, and hope they don't bias the posterior much.
        cols = self.generateRVIndices()
        k = theta[:,cols['k']].flatten()
        t1 = theta[:,cols['t1']].flatten()
        t2 = theta[:,cols['t2']].flatten()
        m1 = k==0
        m2 = k==1
        llh = np.zeros(k.shape[0])
        mk1cov = 1.
        mk2cov = torch.tensor([[1.,.99],[.99,1.]]) 
        if True:
            tf1  = n.transforms.InverseTransform(SinArcSinhTransform(ep[0],dp[0]))
            tf2  = n.transforms.InverseTransform(n.transforms.CompositeTransform([LTransform(torch.cholesky(torch.tensor([[1.,.99],[.99,1.]]))),SAS2DTransform([ep[1],ep[2]],[dp[1],dp[2]])]))
            #tf1  = n.transforms.InverseTransform(L1DTransform(np.sqrt(mk1cov)))
            #tf1  = n.transforms.IdentityTransform()
            #tf2  = n.transforms.InverseTransform(LTransform(torch.cholesky(mk2cov)))
            f1 = Flow(tf1,StandardNormal((1,)))
            f2 = Flow(tf2,StandardNormal((2,)))
            #print(torch.Tensor([t1[m1]]))
            #print(torch.Tensor(np.column_stack([t1[m2],t2[m2]])))
            if m1.sum() > 0:
                #print(t1[m1])
                #print(torch.Tensor(t1[m1].reshape((t1[m1].shape[0],1))).shape)
                #print(torch.Tensor(t1[m1]).shape)
                #print(f1.log_prob(torch.Tensor(t1[m1].reshape((t1[m1].shape[0],1)))).shape)
                #sys.exit(0)
                #print(f1.log_prob(torch.Tensor(t1[m1].reshape((t1[m1].shape[0],1)))).shape)
                llh[m1] = np.log(m1prob) + f1.log_prob(torch.Tensor(t1[m1].reshape((t1[m1].shape[0],1)))).detach().numpy().flatten()
            if m2.sum() > 0:
                #print(torch.Tensor(np.column_stack([t1[m2],t2[m2]])))
                llh[m2] = np.log(1-m1prob) + f2.log_prob(torch.Tensor(np.column_stack([t1[m2],t2[m2]]))).detach().numpy().flatten()
        return llh
//  [python]
def compute_llh(self,theta):
        global ep, dp, m1prob
        # we use the RJ target from Andrieu et al 2009
        # \pi(\theta,k) = 0.25 * N(\theta;0,1) * I(k==1) + 0.75 * N(\theta;[0,0],[[1,-0.9],[-0.9,1]]) * I(k==2)
        # k \in {1,2}
        # What kind of priors do I use here? Seems they are ill-formed. Perhaps really wide uniform priors, and hope they don't bias the posterior much.
        cols = self.generateRVIndices()
        k = theta[:,cols['k']].flatten()
        t1 = theta[:,cols['t1']].flatten()
        t2 = theta[:,cols['t2']].flatten()
        m1 = k==0
        m2 = k==1
        llh = np.zeros(k.shape[0])
        mk1cov = 1.
        mk2cov = torch.tensor([[1.,.99],[.99,1.]]) 
        if True:
            tf1  = n.transforms.InverseTransform(SinArcSinhTransform(ep[0],dp[0]))
            tf2  = n.transforms.InverseTransform(n.transforms.CompositeTransform([LTransform(torch.linalg.cholesky(torch.tensor([[1.,.99],[.99,1.]]))),SAS2DTransform([ep[1],ep[2]],[dp[1],dp[2]])]))
            #tf1  = n.transforms.InverseTransform(L1DTransform(np.sqrt(mk1cov)))
            #tf1  = n.transforms.IdentityTransform()
            #tf2  = n.transforms.InverseTransform(LTransform(torch.linalg.cholesky(mk2cov)))
            f1 = Flow(tf1,StandardNormal((1,)))
            f2 = Flow(tf2,StandardNormal((2,)))
            #print(torch.Tensor([t1[m1]]))
            #print(torch.Tensor(np.column_stack([t1[m2],t2[m2]])))
            if m1.sum() > 0:
                #print(t1[m1])
                #print(torch.Tensor(t1[m1].reshape((t1[m1].shape[0],1))).shape)
                #print(torch.Tensor(t1[m1]).shape)
                #print(f1.log_prob(torch.Tensor(t1[m1]).reshape((-1,1))).shape)
                #sys.exit(0)
                #print(f1.log_prob(torch.Tensor(t1[m1].reshape((t1[m1].shape[0],1)))).shape)
                llh[m1] = np.log(m1prob) + f1.log_prob(torch.Tensor(t1[m1].reshape((t1[m1].shape[0],1)))).detach().numpy().flatten()
            if m2.sum() > 0:
                #print(torch.Tensor(np.column_stack([t1[m2],t2[m2]])))
                llh[m2] = np.log(1-m1prob) + f2.log_prob(torch.Tensor(np.column_stack([t1[m2],t2[m2]]))).detach().numpy().flatten()
        return llh
//  [python]
def _calculate_log_det(self, var):
        """
        Calculate log determinant in log space, to prevent overflow errors.
        args:
            var:            torch.Tensor (1, k, d, d)
        """
        log_det = torch.empty(size=(self.n_components,)).to(var.device)

        for k in range(self.n_components):
            log_det[k] = 2 * torch.log(torch.diagonal(torch.cholesky(var[0, k]))).sum()

        return log_det.unsqueeze(-1)
//  [python]
def _calculate_log_det(self, var):
        """
        Calculate log determinant in log space, to prevent overflow errors.
        args:
            var:            torch.Tensor (1, k, d, d)
        """
        log_det = torch.empty(size=(self.n_components,)).to(var.device)

        for k in range(self.n_components):
            log_det[k] = 2 * torch.log(torch.diagonal(torch.linalg.cholesky(var[0, k]))).sum()

        return log_det.unsqueeze(-1)
//  [python]
def update_covariance(self):
        if not self.is_fit:
            # The precision matrix is positive definite and so we can use its cholesky decomposition to more
            # efficiently compute its inverse (when num_inducing is large)
            try:
                L = torch.cholesky(self.precision)
                self.covariance = Parameter(self.ridge_penalty * L.cholesky_inverse(), requires_grad=False)
                self.is_fit = torch.tensor(True)
                print('cholesky')
            except:
                self.covariance = Parameter(self.ridge_penalty * self.precision.cholesky_inverse(), requires_grad=False)
                self.is_fit = torch.tensor(True)
                print('standard inversion')
        else:
            print(f'no inversion, already fit')
//  [python]
def update_covariance(self):
        if not self.is_fit:
            # The precision matrix is positive definite and so we can use its cholesky decomposition to more
            # efficiently compute its inverse (when num_inducing is large)
            try:
                L = torch.linalg.cholesky(self.precision)
                self.covariance = Parameter(self.ridge_penalty * L.cholesky_inverse(), requires_grad=False)
                self.is_fit = torch.tensor(True)
                print('cholesky')
            except:
                self.covariance = Parameter(self.ridge_penalty * self.precision.cholesky_inverse(), requires_grad=False)
                self.is_fit = torch.tensor(True)
                print('standard inversion')
        else:
            print(f'no inversion, already fit')
//  [python]
def _stable_cholesky(input, upper=False, out=None, eps=1e-6):
    try:
        if upper:
            return torch.cholesky(input, upper=True, out=out).mH
        return torch.cholesky(input, upper=False, out=out)
    except RuntimeError:
        input = condition_scm(input, eps)
        if upper:
            return torch.cholesky(input, upper=True, out=out).mH
        return torch.cholesky(input, upper=False, out=out)
//  [python]
def _stable_cholesky(input, upper=False, out=None, eps=1e-6):
    try:
        if upper:
            return torch.linalg.cholesky(input, out=out).mH
        return torch.linalg.cholesky(input, out=out)
    except RuntimeError:
        input = condition_scm(input, eps)
        if upper:
            return torch.linalg.cholesky(input, out=out).mH
        return torch.linalg.cholesky(input, out=out)
//  [python]
def forward(self, x, L, mass, evals, evecs):

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

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

        if self.method == 'spectral':

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

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

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

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

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

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

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

        else:
            raise ValueError("unrecognized method")

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

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

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

        if self.method == 'spectral':

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

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

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

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

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

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

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

        else:
            raise ValueError("unrecognized method")

        return x_diffuse
//  [python]
def model(self):
        self.set_mode("model")

        N = self.X.size(0)
        Kff = self.kernel(self.X)
        if self.noise.dim() <=1:
            Kff.view(-1)[:: N + 1] += self.jitter + self.noise
        elif self.noise.dim() ==2:
            Kff = Kff + self.noise
            Kff.view(-1)[:: N + 1] += self.jitter
        Lff = torch.cholesky(Kff)

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

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

        zero_loc = self.X.new_zeros(self.X.size(0))
        f_loc = zero_loc + self.mean_function(self.X)
        if self.y is None:
            f_var = Lff.pow(2).sum(dim=-1)
            return f_loc, f_var
        else:
            return pyro.sample(
                self._pyro_get_fullname("y"),
                dist.MultivariateNormal(f_loc, scale_tril=Lff)
                .expand_by(self.y.shape[:-1])
                .to_event(self.y.dim() - 1),
                obs=self.y,
            )
//  [python]
def negative_log_likelihood(self, x, y):
        y_num, y_dimension = y.shape
        Sigma = self.kernel(x, x) + self.log_beta.exp().pow(-1) * torch.eye(
            x.size(0)) + JITTER * torch.eye(x.size(0))

        L = torch.cholesky(Sigma)
        #option 1 (use this if torch supports)
        Gamma,_ = torch.triangular_solve(y, L, upper = False)
        #option 2
        # gamma = L.inverse() @ Y       # we can use this as an alternative because L is a lower triangular matrix.

        nll =  0.5 * (Gamma ** 2).sum() +  L.diag().log().sum() * y_dimension  \
            + 0.5 * y_num * torch.log(2 * torch.tensor(PI)) * y_dimension
        return nll
//  [python]
def negative_log_likelihood(self, x, y):
        y_num, y_dimension = y.shape
        Sigma = self.kernel(x, x) + self.log_beta.exp().pow(-1) * torch.eye(
            x.size(0)) + JITTER * torch.eye(x.size(0))

        L = torch.linalg.cholesky(Sigma)
        #option 1 (use this if torch supports)
        Gamma,_ = torch.triangular_solve(y, L, upper = False)
        #option 2
        # gamma = L.inverse() @ Y       # we can use this as an alternative because L is a lower triangular matrix.

        nll =  0.5 * (Gamma ** 2).sum() +  L.diag().log().sum() * y_dimension  \
            + 0.5 * y_num * torch.log(2 * torch.tensor(PI)) * y_dimension
        return nll
//  [python]
def model(self):
        self.set_mode("model")
        N = self.X.size(0)
        x_noise = pyro.sample('xerr',dist.Normal(torch.zeros(N),self.x_noise**0.5).to_event(1))
        if self.X.dim()<=1:
            X_noisy = self.X
            X_noisy = (self.X+x_noise)
        else:
            X_noisy = torch.clone(self.X)
            X_noisy[:,0] += x_noise
       
        Kff = self.kernel(X_noisy).contiguous()
        Kff = Kff + self.noise
        Kff.view(-1)[:: N + 1] += self.jitter
        # Kff.view(-1)[:: N + 1] += self.jitter + self.noise  # add noise to diagonal
        Lff = torch.cholesky(Kff)
        zero_loc = X_noisy.new_zeros(X_noisy.size(0))
        f_loc = zero_loc + self.mean_function(X_noisy)
        if self.y is None:
            f_var = Lff.pow(2).sum(dim=-1)
            return f_loc, f_var
        else:

            return pyro.sample(
                self._pyro_get_fullname("y"),
                dist.MultivariateNormal(f_loc, scale_tril=Lff)
                .expand_by(self.y.shape[:-1])
                .to_event(self.y.dim() - 1),
                obs=self.y,
            )