jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def sample(self, X_in: torch.Tensor, **kwargs) -> torch.Tensor:
        """Sample current layer neuron values given upper layer neuron values. When
        self._sample_strat is LayerSampleStrat.MAP and LayerSampleStrat.P_PRED, sample from
        p(X_out| argmax_W p(W), self._Sigma, ...) and E[p(X_out| W, self._Sigma, ...)].
        NOTE: Treats elements in batch independently.

        Args:
            X_in (torch.Tensor): Upper layer neuron values of shape <d_batch x d_in>.

        Returns:
            torch.Tensor: Sampled lower layer neuron values of shape <d_batch x d_out>.
        """
        if self._update_strat == LayerUpdateStrat.KERNEL:
            # HACK: Create a new class for kernel stuff later
            mus, Sigmas = kernel_posterior_params(X_train=self._X_train, Y_train=self._Y_train,
                                                  X_test=X_in, kernel=self._kernel_type,
                                                  kernel_params=self._kernel_params,
                                                  Sigma_trtr_inv=self._X_train_kernel)
            L = torch.linalg.cholesky(Sigmas)
            return mus + L.matmul(torch.randn_like(mus))

        marginal_mean = self.predict(X_in=X_in)
        if self._sample_strat == LayerSampleStrat.MAP:
            marginal_Sigma = self._Sigma
        elif self._sample_strat == LayerSampleStrat.P_PRED:
            Z_in = self._f(X_in)
            marginal_Sigma = self._Sigma + Z_in.matmul(self._U).matmul(Z_in.T).diag()
            marginal_Sigma = marginal_Sigma.unsqueeze(-1)
        else:
            raise NotImplementedError()
        dist = pdists.Normal(torch.zeros_like(marginal_mean), marginal_Sigma ** 0.5)
        obs = None if kwargs.get('X_obs') is None else (kwargs['X_obs']-marginal_mean)
        sample = marginal_mean + pyro.sample(self._name, dist, obs=obs)
        return sample
//  [python]
def batch_cholesky_inverse(matrix):
    """
    Args:
        matrix, 4-D tensor, [N, C, M, M]. matrix must be a symmetric positive define matrix.
    """
    chol_low = torch.cholesky(matrix, upper=False)
    chol_low_inv = batch_low_tri_inv(chol_low)
    return torch.matmul(chol_low_inv.transpose(-2, -1), chol_low_inv)
//  [python]
def batch_cholesky_inverse(matrix):
    """
    Args:
        matrix, 4-D tensor, [N, C, M, M]. matrix must be a symmetric positive define matrix.
    """
    chol_low = torch.linalg.cholesky(matrix, upper=False)
    chol_low_inv = batch_low_tri_inv(chol_low)
    return torch.matmul(chol_low_inv.transpose(-2, -1), chol_low_inv)
//  [python]
def log_det_by_cholesky(matrix):
    """
    Args:
        matrix: matrix must be a positive define matrix.shape [N, C, D, D].
    Ref:
       https://github.com/tensorflow/tensorflow/blob/r1.13/tensorflow/python/ops/linalg/linalg_impl.py
    """
    # This uses the property that the log det(A) = 2 * sum(log(real(diag(C))))
    # where C is the cholesky decomposition of A.
    chol = torch.cholesky(matrix)
    # return 2.0 * torch.sum(torch.log(torch.diagonal(chol, dim1=-2, dim2=-1) + 1e-6), dim=-1)
    return 2.0 * torch.sum(torch.log(torch.diagonal(chol, dim1=-2, dim2=-1) + 1e-8), dim=-1)
//  [python]
def log_det_by_cholesky(matrix):
    """
    Args:
        matrix: matrix must be a positive define matrix.shape [N, C, D, D].
    Ref:
       https://github.com/tensorflow/tensorflow/blob/r1.13/tensorflow/python/ops/linalg/linalg_impl.py
    """
    # This uses the property that the log det(A) = 2 * sum(log(real(diag(C))))
    # where C is the cholesky decomposition of A.
    chol = torch.linalg.cholesky(matrix)
    # return 2.0 * torch.sum(torch.log(torch.diagonal(chol, dim1=-2, dim2=-1) + 1e-6), dim=-1)
    return 2.0 * torch.sum(torch.log(torch.diagonal(chol, dim1=-2, dim2=-1) + 1e-8), dim=-1)
//  [python]
def mahalanobis_dist(diff: torch.Tensor, covariance: torch.Tensor) -> torch.Tensor:
    cholesky = torch.cholesky(covariance)
    y = torch.cholesky_solve(diff.unsqueeze(-1), cholesky).squeeze(-1)
    return torch.sqrt(torch.sum(diff * y, -1))
//  [python]
def mahalanobis_dist(diff: torch.Tensor, covariance: torch.Tensor) -> torch.Tensor:
    cholesky = torch.linalg.cholesky(covariance)
    y = torch.cholesky_solve(diff.unsqueeze(-1), cholesky).squeeze(-1)
    return torch.sqrt(torch.sum(diff * y, -1))
//  [python]
def update_posterior(self):
        self.ppt *= 0
        self.py *= 0
        vark = self.var_k
        if self.total_steps > self.config.exploration_steps:
            batch_size = 32
            num_iters = int(min(self.target_batch_size, self.total_steps)/batch_size)
            for _ in range(num_iters):
                transitions = self.replay.usample(batch_size)
                states = self.config.state_normalizer(transitions.state)
                next_states = self.config.state_normalizer(transitions.next_state)
                masks = tensor(transitions.mask)
                rewards = tensor(transitions.reward)
                actions = tensor(transitions.action).long()

                with torch.no_grad():
                    policy_state_rep, y_target =  self.find_state_rep(states, next_states, masks, rewards, actions)

                # can improve by not using a loop
                for idx in range(batch_size):
                    self.ppt[int(actions[idx]), : , :] += torch.matmul(policy_state_rep[idx].unsqueeze(0).T, policy_state_rep[idx].unsqueeze(0))
                    self.py[int(actions[idx]), :] += policy_state_rep[idx].T * y_target[idx].item()

            for idx in range(self.num_actions):
                inv = torch.inverse(self.ppt[idx]/self.noise_var + 1/self.prior_var*tensor(torch.eye(self.layer_size)))
                self.policy_mean[idx] = torch.matmul(inv, self.py[idx])/self.noise_var
                self.policy_cov[idx] = vark * inv
                try:
                    self.cov_decom[idx] = torch.cholesky((self.policy_cov[idx]+self.policy_cov[idx].T)/2)
                except RuntimeError:
                    pass
//  [python]
def update_posterior(self):
        self.ppt *= 0
        self.py *= 0
        vark = self.var_k
        if self.total_steps > self.config.exploration_steps:
            batch_size = 32
            num_iters = int(min(self.target_batch_size, self.total_steps)/batch_size)
            for _ in range(num_iters):
                transitions = self.replay.usample(batch_size)
                states = self.config.state_normalizer(transitions.state)
                next_states = self.config.state_normalizer(transitions.next_state)
                masks = tensor(transitions.mask)
                rewards = tensor(transitions.reward)
                actions = tensor(transitions.action).long()

                with torch.no_grad():
                    policy_state_rep, y_target =  self.find_state_rep(states, next_states, masks, rewards, actions)

                # can improve by not using a loop
                for idx in range(batch_size):
                    self.ppt[int(actions[idx]), : , :] += torch.matmul(policy_state_rep[idx].unsqueeze(0).T, policy_state_rep[idx].unsqueeze(0))
                    self.py[int(actions[idx]), :] += policy_state_rep[idx].T * y_target[idx].item()

            for idx in range(self.num_actions):
                inv = torch.inverse(self.ppt[idx]/self.noise_var + 1/self.prior_var*tensor(torch.eye(self.layer_size)))
                self.policy_mean[idx] = torch.matmul(inv, self.py[idx])/self.noise_var
                self.policy_cov[idx] = vark * inv
                try:
                    self.cov_decom[idx] = torch.linalg.cholesky((self.policy_cov[idx]+self.policy_cov[idx].T)/2)
                except RuntimeError:
                    pass
//  [python]
def matrix_parallel_transport(self, sigma, g, h):
        Lp = torch.cholesky(sigma).transpose(-2, -1)
        A_g_h_L = self.parallel_transport(Lp, g, h)
        sigma_parallel = A_g_h_L.transpose(-2, -1) @ A_g_h_L
        return sigma_parallel
//  [python]
def matrix_parallel_transport(self, sigma, g, h):
        Lp = torch.linalg.cholesky(sigma).transpose(-2, -1)
        A_g_h_L = self.parallel_transport(Lp, g, h)
        sigma_parallel = A_g_h_L.transpose(-2, -1) @ A_g_h_L
        return sigma_parallel
//  [python]
def __init__(self, config):
        print('Number of actions:', config.action_dim)
        BaseAgent.__init__(self, config)
        self.config = config
        config.lock = mp.Lock()

        self.replay = config.replay_fn()
        self.network = config.network_fn()
        self.network.share_memory()
        self.optimizer = config.optimizer_fn(self.network.parameters())
        self.target_network = config.network_fn()

        self.actor = BDQNActor(config, self.logger)
        self.target_batch_size = 200000
        self.prior_var = config.prior_var
        self.noise_var = config.noise_var
        self.var_k = config.var_k
        self.num_actions = config.action_dim
        self.layer_size = 512
        self.sampled_mean = tensor(torch.normal(0, 0.01, size=(self.num_actions, self.layer_size)))
        self.policy_mean = tensor(torch.normal(0, 0.01, size=(self.num_actions, self.layer_size)))
        self.policy_cov = tensor(torch.normal(0, 1, size=(self.num_actions, self.layer_size, self.layer_size)))
        self.cov_decom = self.policy_cov
        for idx in range(self.num_actions):
            self.policy_cov[idx] = torch.eye(self.layer_size)
            self.cov_decom[idx] = torch.cholesky((self.policy_cov[idx] + self.policy_cov[idx].T)/2.0)
        self.ppt = tensor(torch.zeros(self.num_actions, self.layer_size, self.layer_size))
        self.py = tensor(torch.zeros(self.num_actions, self.layer_size))
        self.bayes_reg = False
        self.actor.update_mean(self.sampled_mean)
        self.target_network.load_state_dict(self.network.state_dict())
        self.actor.set_network(self.network)
        self.total_steps = 0
        self.episode_length = 0
        self.num_episodes = 0
//  [python]
def __init__(self, config):
        print('Number of actions:', config.action_dim)
        BaseAgent.__init__(self, config)
        self.config = config
        config.lock = mp.Lock()

        self.replay = config.replay_fn()
        self.network = config.network_fn()
        self.network.share_memory()
        self.optimizer = config.optimizer_fn(self.network.parameters())
        self.target_network = config.network_fn()

        self.actor = BDQNActor(config, self.logger)
        self.target_batch_size = 200000
        self.prior_var = config.prior_var
        self.noise_var = config.noise_var
        self.var_k = config.var_k
        self.num_actions = config.action_dim
        self.layer_size = 512
        self.sampled_mean = tensor(torch.normal(0, 0.01, size=(self.num_actions, self.layer_size)))
        self.policy_mean = tensor(torch.normal(0, 0.01, size=(self.num_actions, self.layer_size)))
        self.policy_cov = tensor(torch.normal(0, 1, size=(self.num_actions, self.layer_size, self.layer_size)))
        self.cov_decom = self.policy_cov
        for idx in range(self.num_actions):
            self.policy_cov[idx] = torch.eye(self.layer_size)
            self.cov_decom[idx] = torch.linalg.cholesky((self.policy_cov[idx] + self.policy_cov[idx].T)/2.0)
        self.ppt = tensor(torch.zeros(self.num_actions, self.layer_size, self.layer_size))
        self.py = tensor(torch.zeros(self.num_actions, self.layer_size))
        self.bayes_reg = False
        self.actor.update_mean(self.sampled_mean)
        self.target_network.load_state_dict(self.network.state_dict())
        self.actor.set_network(self.network)
        self.total_steps = 0
        self.episode_length = 0
        self.num_episodes = 0
//  [python]
def log_det(self, A):
        assert A.dim() in [2, 3]
        # regularize when computing log-det
        A = A + self.cfg.design_cov_reg * torch.eye(A.shape[1], device=A.device)
        return 2 * torch.cholesky(A).diagonal(dim1=-2, dim2=-1).log().sum(-1)
//  [python]
def log_det(self, A):
        assert A.dim() in [2, 3]
        # regularize when computing log-det
        A = A + self.cfg.design_cov_reg * torch.eye(A.shape[1], device=A.device)
        return 2 * torch.linalg.cholesky(A).diagonal(dim1=-2, dim2=-1).log().sum(-1)
//  [python]
def cholesky_eigh(A, B, reg_B=1e-6, n_eig=None):
    """
    -- Generalized eigenvalue problem: A * v_i = lambda_i * B * v_i --

    First apply cholesky decomposition to B and then solve the generalized eigvalue problem.

    Notes
    -----
    The eigenvecs object which is returned is a matrix whose column eigvecs[:,i] is the eigenvector associated to eigvals[i]
    """

    # check that both matrices are symmetric
    if not torch.allclose(A.transpose(0, 1), A):
        raise ValueError(
            "The matrices need to be symmetric to solve the generalized eigenvalue problem via cholesky decomposition. A >> ",
            A,
        )
    if not (torch.allclose(B.transpose(0, 1), B)):
        raise ValueError(
            "The matrices need to be symmetric to solve the generalized eigenvalue problem via cholesky decomposition. A >> ",
            B,
        )

    # (0) regularize B matrix before cholesky
    if reg_B is not None:
        B = B + reg_B * torch.eye(B.shape[0]).to(B.device)

    # (1) use cholesky decomposition for B
    L = torch.cholesky(B, upper=False)

    # (2) define new matrix using cholesky decomposition
    L_t = torch.t(L)
    L_ti = torch.inverse(L_t)
    L_i = torch.inverse(L)
    A_new = torch.matmul(torch.matmul(L_i, A), L_ti)

    # (3) find eigenvalues and vectors of A_new
    eigvals, eigvecs = torch.linalg.eigh(A_new, UPLO="L")
    # sort
    eigvals, indices = torch.sort(eigvals, 0, descending=True)
    eigvecs = eigvecs[:, indices]

    # (4) return to original eigenvectors
    eigvecs = torch.matmul(L_ti, eigvecs)

    # (5) normalize them
    eigvecs = torch.nn.functional.normalize(eigvecs, dim=0)
    # set the first component positive
    eigvecs = eigvecs.mul(torch.sign(eigvecs[0, :]).unsqueeze(0).expand_as(eigvecs))

    # (6) keep only first n_eig eigvals and eigvecs
    if n_eig is not None:
        eigvals = eigvals[:n_eig]
        eigvecs = eigvecs[:, :n_eig]

    return eigvals, eigvecs
//  [python]
def cholesky_eigh(A, B, reg_B=1e-6, n_eig=None):
    """
    -- Generalized eigenvalue problem: A * v_i = lambda_i * B * v_i --

    First apply cholesky decomposition to B and then solve the generalized eigvalue problem.

    Notes
    -----
    The eigenvecs object which is returned is a matrix whose column eigvecs[:,i] is the eigenvector associated to eigvals[i]
    """

    # check that both matrices are symmetric
    if not torch.allclose(A.transpose(0, 1), A):
        raise ValueError(
            "The matrices need to be symmetric to solve the generalized eigenvalue problem via cholesky decomposition. A >> ",
            A,
        )
    if not (torch.allclose(B.transpose(0, 1), B)):
        raise ValueError(
            "The matrices need to be symmetric to solve the generalized eigenvalue problem via cholesky decomposition. A >> ",
            B,
        )

    # (0) regularize B matrix before cholesky
    if reg_B is not None:
        B = B + reg_B * torch.eye(B.shape[0]).to(B.device)

    # (1) use cholesky decomposition for B
    L = torch.linalg.cholesky(B, upper=False)

    # (2) define new matrix using cholesky decomposition
    L_t = torch.t(L)
    L_ti = torch.inverse(L_t)
    L_i = torch.inverse(L)
    A_new = torch.matmul(torch.matmul(L_i, A), L_ti)

    # (3) find eigenvalues and vectors of A_new
    eigvals, eigvecs = torch.linalg.eigh(A_new, UPLO="L")
    # sort
    eigvals, indices = torch.sort(eigvals, 0, descending=True)
    eigvecs = eigvecs[:, indices]

    # (4) return to original eigenvectors
    eigvecs = torch.matmul(L_ti, eigvecs)

    # (5) normalize them
    eigvecs = torch.nn.functional.normalize(eigvecs, dim=0)
    # set the first component positive
    eigvecs = eigvecs.mul(torch.sign(eigvecs[0, :]).unsqueeze(0).expand_as(eigvecs))

    # (6) keep only first n_eig eigvals and eigvecs
    if n_eig is not None:
        eigvals = eigvals[:n_eig]
        eigvecs = eigvecs[:, :n_eig]

    return eigvals, eigvecs
//  [python]
def setup_item_vals(self):
        self.N_x,self.m = self.X.shape
        self.reg = self.N_x*self.lamb
        self.eye = torch.eye(self.N_x).to(self.device)*self.reg
        L = torch.cholesky(self.k(self.X)+ self.eye)
        self.precond = torch.cholesky_inverse(L)
        self.setup_the_rest()
//  [python]
def setup_item_vals(self):
        self.N_x,self.m = self.X.shape
        self.reg = self.N_x*self.lamb
        self.eye = torch.eye(self.N_x).to(self.device)*self.reg
        L = torch.linalg.cholesky(self.k(self.X)+ self.eye)
        self.precond = torch.cholesky_inverse(L)
        self.setup_the_rest()
//  [python]
def get_weight(self):
        Sigma = self.kernel @ self.kernel.T
        eps = Sigma.diag().mean().div(1000.).item()
        Sigma = Sigma + eps * torch.eye(
            Sigma.shape[0], device=Sigma.device, dtype=Sigma.dtype)
        L = torch.cholesky(Sigma)
        kernel_ = torch.triangular_solve(self.kernel, L, upper=False)[0]
        kernel_ = kernel_.reshape(self.feat_dim, self.feat_dim, 4, 4)

        Sigma = self.weight @ self.weight.T
        eps = Sigma.diag().mean().div(1000.).item()
        Sigma = Sigma + eps * torch.eye(
            Sigma.shape[0], device=Sigma.device, dtype=Sigma.dtype)
        L = torch.cholesky(Sigma)
        weight_ = torch.triangular_solve(self.weight, L, upper=False)[0]
        return kernel_, weight_
//  [python]
def get_weight(self):
        Sigma = self.kernel @ self.kernel.T
        eps = Sigma.diag().mean().div(1000.).item()
        Sigma = Sigma + eps * torch.eye(
            Sigma.shape[0], device=Sigma.device, dtype=Sigma.dtype)
        L = torch.linalg.cholesky(Sigma)
        kernel_ = torch.linalg.solve_triangular(L, self.kernel, upper=False)
        kernel_ = kernel_.reshape(self.feat_dim, self.feat_dim, 4, 4)

        Sigma = self.weight @ self.weight.T
        eps = Sigma.diag().mean().div(1000.).item()
        Sigma = Sigma + eps * torch.eye(
            Sigma.shape[0], device=Sigma.device, dtype=Sigma.dtype)
        L = torch.linalg.cholesky(Sigma)
        weight_ = torch.linalg.solve_triangular(L, self.weight, upper=False)
        return kernel_, weight_
//  [python]
def update_posteriors(self, dataset: Dataset):
        """
        Compute feature maps for all states in the replay buffer, then update posteriors with Bayesian Linear
        Regression. Cholesky decomposition for computing the inverse is used with additional checks to ensure
        positive-semi definite matrices.
        """

        x, y = self.compute_feature_maps(dataset)
        Phi_pre = x.T.matmul(x) * self.noise_variance
        coeff = BLR_COEFFICIENT

        while True:
            input_matrix = Phi_pre.double() + (self.eye.double() * coeff)
            try:
                chol = torch.cholesky(input_matrix)
                Phi = torch.cholesky_inverse(chol).float()
            except:
                coeff *= 10
                continue

            if Phi.isnan().any():
                coeff *= 10
                continue

            self.reward_cov = Phi * self.reward_prior
            self.transition_cov = Phi * self.transition_prior
            break
            
        for i in range(self.embed_dim + 1):
            self.mu[i] = (self.noise_variance * Phi).matmul(x.T.matmul(y[:, i]))
//  [python]
def update_posteriors(self, dataset: Dataset):
        """
        Compute feature maps for all states in the replay buffer, then update posteriors with Bayesian Linear
        Regression. Cholesky decomposition for computing the inverse is used with additional checks to ensure
        positive-semi definite matrices.
        """

        x, y = self.compute_feature_maps(dataset)
        Phi_pre = x.T.matmul(x) * self.noise_variance
        coeff = BLR_COEFFICIENT

        while True:
            input_matrix = Phi_pre.double() + (self.eye.double() * coeff)
            try:
                chol = torch.linalg.cholesky(input_matrix)
                Phi = torch.cholesky_inverse(chol).float()
            except:
                coeff *= 10
                continue

            if Phi.isnan().any():
                coeff *= 10
                continue

            self.reward_cov = Phi * self.reward_prior
            self.transition_cov = Phi * self.transition_prior
            break
            
        for i in range(self.embed_dim + 1):
            self.mu[i] = (self.noise_variance * Phi).matmul(x.T.matmul(y[:, i]))
//  [python]
def safe_cholesky(A, epsilon=1.0e-8, num_tries=5):
    """
    Equivalent of torch.linalg.cholesky that progressively adds
    diagonal jitter to avoid cholesky errors.
    """
    try:
        return torch.cholesky(A)
    except RuntimeError as e:
        Aprime = A.clone()
        jitter_prev = 0.0
        for i in range(num_tries):
            jitter_new = epsilon * (10 ** i)
            Aprime.diagonal(dim1=-2, dim2=-1).add_(jitter_new - jitter_prev)
            jitter_prev = jitter_new
            try:
                return torch.cholesky(Aprime)
            except RuntimeError:
                continue
        raise e
//  [python]
def safe_cholesky(A, epsilon=1.0e-8, num_tries=5):
    """
    Equivalent of torch.linalg.cholesky that progressively adds
    diagonal jitter to avoid cholesky errors.
    """
    try:
        return torch.linalg.cholesky(A)
    except RuntimeError as e:
        Aprime = A.clone()
        jitter_prev = 0.0
        for i in range(num_tries):
            jitter_new = epsilon * (10 ** i)
            Aprime.diagonal(dim1=-2, dim2=-1).add_(jitter_new - jitter_prev)
            jitter_prev = jitter_new
            try:
                return torch.linalg.cholesky(Aprime)
            except RuntimeError:
                continue
        raise e
//  [python]
def calculate_vel_impulse(rela_v, J_c, elasticity):

    # set up differentiable optimization problem
    f = cp.Variable((2, 1)) # [f_t, f_n]
    A_decom_p = cp.Parameter((2, 2))
    v_p = cp.Parameter((2, 1))
    objective = cp.Minimize(0.5 * cp.sum_squares(A_decom_p @ f) + cp.sum(cp.multiply(f, v_p)))
    constraints = [
        f[0] >= 0, # normal impulse should be creater than zero
        f[1] == 0,
    ]
    problem = cp.Problem(objective, constraints)
    cvxpylayer = CvxpyLayer(problem, parameters=[A_decom_p, v_p], variables=[f])

    # feed values to the layer
    A = J_c.transpose(0, 1) @ J_c
    A_decom = torch.cholesky(A)
    impulse, = cvxpylayer(
        A_decom,
        J_c @ rela_v.reshape(-1, 1),
    )
    return (J_c.transpose(0, 1) @ impulse)[:, 0] * (1+elasticity)
//  [python]
def calculate_vel_impulse(rela_v, J_c, elasticity):

    # set up differentiable optimization problem
    f = cp.Variable((2, 1)) # [f_t, f_n]
    A_decom_p = cp.Parameter((2, 2))
    v_p = cp.Parameter((2, 1))
    objective = cp.Minimize(0.5 * cp.sum_squares(A_decom_p @ f) + cp.sum(cp.multiply(f, v_p)))
    constraints = [
        f[0] >= 0, # normal impulse should be creater than zero
        f[1] == 0,
    ]
    problem = cp.Problem(objective, constraints)
    cvxpylayer = CvxpyLayer(problem, parameters=[A_decom_p, v_p], variables=[f])

    # feed values to the layer
    A = J_c.transpose(0, 1) @ J_c
    A_decom = torch.linalg.cholesky(A)
    impulse, = cvxpylayer(
        A_decom,
        J_c @ rela_v.reshape(-1, 1),
    )
    return (J_c.transpose(0, 1) @ impulse)[:, 0] * (1+elasticity)
//  [python]
def predict(gamma, Z, labels):
    Y = torch.zeros(n_runs,n_lsamples, n_ways,device='cuda')
    Y.scatter_(2,labels[:,:n_lsamples].unsqueeze(2), 1)
    delta = torch.sum(Z, 1)
    iden = torch.eye(5,device='cuda')
    iden = iden.reshape((1, 5, 5))
    iden = iden.repeat(10000, 1, 1)
    W = torch.bmm(torch.transpose(Z,1,2), Z/delta.unsqueeze(1))
    L = iden - W
    Z_l = Z[:,:n_lsamples]
    u = torch.cholesky(torch.bmm(torch.transpose(Z_l,1,2), Z_l) + gamma * L)# + 0.1*iden)
    A = torch.cholesky_solve(torch.bmm(torch.transpose(Z_l,1,2), Y), u)
    Pred = Z.bmm(A)
    normalizer = torch.sum(Pred,dim=1,keepdim=True)
    # #normalizer = Pred[:,:n_lsamples].max(dim=1)[0].unsqueeze(1)
    Pred = (n_shot+n_queries)*Pred/normalizer
    # normalizer = torch.sum(Pred, dim=2, keepdim=True)
    # Pred = Pred/normalizer
    # Pred[:, :n_lsamples].fill_(0)
    # Pred[:, :n_lsamples].scatter_(2, labels[:, :n_lsamples].unsqueeze(2), 1)
    # N = PredZ.shape[0]
    # K = PredZ.shape[1]
    # pred = np.zeros((N, K))
    #
    # for k in range(K):
    #     current_pred = np.dot(Z, A[:, k])

    return Pred
//  [python]
def predict(gamma, Z, labels):
    Y = torch.zeros(n_runs,n_lsamples, n_ways,device='cuda')
    Y.scatter_(2,labels[:,:n_lsamples].unsqueeze(2), 1)
    delta = torch.sum(Z, 1)
    iden = torch.eye(5,device='cuda')
    iden = iden.reshape((1, 5, 5))
    iden = iden.repeat(10000, 1, 1)
    W = torch.bmm(torch.transpose(Z,1,2), Z/delta.unsqueeze(1))
    L = iden - W
    Z_l = Z[:,:n_lsamples]
    u = torch.linalg.cholesky(torch.bmm(torch.transpose(Z_l,1,2), Z_l) + gamma * L)# + 0.1*iden)
    A = torch.cholesky_solve(torch.bmm(torch.transpose(Z_l,1,2), Y), u)
    Pred = Z.bmm(A)
    normalizer = torch.sum(Pred,dim=1,keepdim=True)
    # #normalizer = Pred[:,:n_lsamples].max(dim=1)[0].unsqueeze(1)
    Pred = (n_shot+n_queries)*Pred/normalizer
    # normalizer = torch.sum(Pred, dim=2, keepdim=True)
    # Pred = Pred/normalizer
    # Pred[:, :n_lsamples].fill_(0)
    # Pred[:, :n_lsamples].scatter_(2, labels[:, :n_lsamples].unsqueeze(2), 1)
    # N = PredZ.shape[0]
    # K = PredZ.shape[1]
    # pred = np.zeros((N, K))
    #
    # for k in range(K):
    #     current_pred = np.dot(Z, A[:, k])

    return Pred
//  [python]
def __init__(self, file_name, num_bins_per_dim, use_whitened):
        self.use_whitened = use_whitened
        self._num_latents = num_bins_per_dim ** 2
        self._num_grid_per_dim = num_bins_per_dim
        bin_counts = np.array(get_bin_counts(read_points(file_name), num_bins_per_dim))

        self._flat_bin_counts = bin_counts.flatten()

        self._poisson_a = 1.0 / self._num_latents

        self._signal_variance = 1.91
        self._beta = 1.0 / 33

        # torch
        self._bin_vals = th.from_numpy(get_bin_vals(num_bins_per_dim)).cuda().float()

        short_kernel_func = lambda x, y: th_batch_kernel_fn(
            x, y, self._signal_variance, num_bins_per_dim, self._beta
        )
        self._gram_matrix = short_kernel_func(self._bin_vals, self._bin_vals)
        self._cholesky_gram = th.cholesky(self._gram_matrix)
        self._white_gaussian_log_normalizer = (
            -0.5 * self._num_latents * np.log(2.0 * np.pi)
        )  # float

        half_log_det_gram = th.sum(th.log(th.abs(th.diag(self._cholesky_gram))))
        self._unwhitened_gaussian_log_normalizer = (
            -0.5 * self._num_latents * np.log(2.0 * np.pi) - half_log_det_gram
        )  # tensor scalar

        self._mu_zero = np.log(126.0) - 0.5 * self._signal_variance  # tensor scalar

        self._flat_bin_counts = th.from_numpy(self._flat_bin_counts).cuda().float()

        if use_whitened:
            self.evaluate_log_density = self.whitened_posterior_log_density
        else:
            self.evaluate_log_density = self.unwhitened_posterior_log_density
//  [python]
def __init__(self, file_name, num_bins_per_dim, use_whitened):
        self.use_whitened = use_whitened
        self._num_latents = num_bins_per_dim ** 2
        self._num_grid_per_dim = num_bins_per_dim
        bin_counts = np.array(get_bin_counts(read_points(file_name), num_bins_per_dim))

        self._flat_bin_counts = bin_counts.flatten()

        self._poisson_a = 1.0 / self._num_latents

        self._signal_variance = 1.91
        self._beta = 1.0 / 33

        # torch
        self._bin_vals = th.from_numpy(get_bin_vals(num_bins_per_dim)).cuda().float()

        short_kernel_func = lambda x, y: th_batch_kernel_fn(
            x, y, self._signal_variance, num_bins_per_dim, self._beta
        )
        self._gram_matrix = short_kernel_func(self._bin_vals, self._bin_vals)
        self._cholesky_gram = th.linalg.cholesky(self._gram_matrix)
        self._white_gaussian_log_normalizer = (
            -0.5 * self._num_latents * np.log(2.0 * np.pi)
        )  # float

        half_log_det_gram = th.sum(th.log(th.abs(th.diag(self._cholesky_gram))))
        self._unwhitened_gaussian_log_normalizer = (
            -0.5 * self._num_latents * np.log(2.0 * np.pi) - half_log_det_gram
        )  # tensor scalar

        self._mu_zero = np.log(126.0) - 0.5 * self._signal_variance  # tensor scalar

        self._flat_bin_counts = th.from_numpy(self._flat_bin_counts).cuda().float()

        if use_whitened:
            self.evaluate_log_density = self.whitened_posterior_log_density
        else:
            self.evaluate_log_density = self.unwhitened_posterior_log_density
//  [python]
def create_tril(scales: torch.Tensor) -> torch.Tensor:
    """Create a 1 dimentional tensor containing a flatten tridiagonal matrix.

    A covariance matrix is created using scales**2 for variances and the
    covariances are set to zero. A tridiagonal is created using the
    cholesky decomposition and the diagonal elements are replaced with
    their log.

    :param scales: standard deviations
    :return:
    """
    dim = len(scales)
    cov = torch.full((dim, dim), 0.0)
    cov[range(dim), range(dim)] = scales * scales
    tril = torch.cholesky(cov)
    indices = torch.tril_indices(row=dim, col=dim, offset=0)
    tril[range(dim), range(dim)] = tril[range(dim), range(dim)].log()
    return tril[indices[0], indices[1]]
//  [python]
def create_tril(scales: torch.Tensor) -> torch.Tensor:
    """Create a 1 dimentional tensor containing a flatten tridiagonal matrix.

    A covariance matrix is created using scales**2 for variances and the
    covariances are set to zero. A tridiagonal is created using the
    cholesky decomposition and the diagonal elements are replaced with
    their log.

    :param scales: standard deviations
    :return:
    """
    dim = len(scales)
    cov = torch.full((dim, dim), 0.0)
    cov[range(dim), range(dim)] = scales * scales
    tril = torch.linalg.cholesky(cov)
    indices = torch.tril_indices(row=dim, col=dim, offset=0)
    tril[range(dim), range(dim)] = tril[range(dim), range(dim)].log()
    return tril[indices[0], indices[1]]
//  [python]
def draw(self, x, y, savestate=False, offset=0):
        proposed = {key : None for key in self.keys}

        if self.counter.num_batches != 1:
            self.current['target_val'] = self.model.log_target(self.current['sample'].clone().detach(), x, y)

        randn_sample = torch.randn(self.model.num_params(), dtype=self.model.dtype, device=self.model.device)
        if (self.counter.idx + 1 - offset > self.t0):
            if torch.rand(1, dtype=self.model.dtype, device=self.model.device) < self.l:
                proposed['sample'] = self.current['sample'].clone().detach() + self.c * randn_sample
            else:
                proposed['sample'] = \
                    self.current['sample'].clone().detach() + self.b * torch.cholesky(self.cov) @ randn_sample
        else:
            proposed['sample'] = self.current['sample'].clone().detach() + self.c * randn_sample
        proposed['target_val'] = self.model.log_target(proposed['sample'].clone().detach(), x, y)

        log_rate = proposed['target_val'] - self.current['target_val']

        if torch.log(torch.rand(1, dtype=self.model.dtype, device=self.model.device)) < log_rate:
            self.current['sample'] = proposed['sample'].clone().detach()
            if self.counter.num_batches == 1:
                self.current['target_val'] = proposed['target_val'].clone().detach()
            self.current['accepted'] = 1
            if (self.counter.idx > 0):
                self.num_accepted = self.num_accepted + 1
        else:
            self.model.set_params(self.current['sample'].clone().detach())
            self.current['accepted'] = 0

        self.running_mean = recursive_mean(
            self.running_mean, self.counter.idx + 1, self.current['sample'], offset=offset
        )
        self.cov_sum = self.cov_sum + torch.ger(self.current['sample'], self.current['sample'])
        if (self.counter.idx + 1 - offset >= self.t0):
            if (self.num_accepted == 0):
                self.cov = self.cov0.clone().detach()
            else:
                self.set_recursive_cov(self.counter.idx, offset=offset)
                if self.transform is not None:
                    self.cov = self.transform(self.cov)

        if savestate:
            self.chain.detach_and_update(self.current)

        self.current['sample'].detach_()
        self.current['target_val'].detach_()
//  [python]
def draw(self, x, y, savestate=False, offset=0):
        proposed = {key : None for key in self.keys}

        if self.counter.num_batches != 1:
            self.current['target_val'] = self.model.log_target(self.current['sample'].clone().detach(), x, y)

        randn_sample = torch.randn(self.model.num_params(), dtype=self.model.dtype, device=self.model.device)
        if (self.counter.idx + 1 - offset > self.t0):
            if torch.rand(1, dtype=self.model.dtype, device=self.model.device) < self.l:
                proposed['sample'] = self.current['sample'].clone().detach() + self.c * randn_sample
            else:
                proposed['sample'] = \
                    self.current['sample'].clone().detach() + self.b * torch.linalg.cholesky(self.cov) @ randn_sample
        else:
            proposed['sample'] = self.current['sample'].clone().detach() + self.c * randn_sample
        proposed['target_val'] = self.model.log_target(proposed['sample'].clone().detach(), x, y)

        log_rate = proposed['target_val'] - self.current['target_val']

        if torch.log(torch.rand(1, dtype=self.model.dtype, device=self.model.device)) < log_rate:
            self.current['sample'] = proposed['sample'].clone().detach()
            if self.counter.num_batches == 1:
                self.current['target_val'] = proposed['target_val'].clone().detach()
            self.current['accepted'] = 1
            if (self.counter.idx > 0):
                self.num_accepted = self.num_accepted + 1
        else:
            self.model.set_params(self.current['sample'].clone().detach())
            self.current['accepted'] = 0

        self.running_mean = recursive_mean(
            self.running_mean, self.counter.idx + 1, self.current['sample'], offset=offset
        )
        self.cov_sum = self.cov_sum + torch.ger(self.current['sample'], self.current['sample'])
        if (self.counter.idx + 1 - offset >= self.t0):
            if (self.num_accepted == 0):
                self.cov = self.cov0.clone().detach()
            else:
                self.set_recursive_cov(self.counter.idx, offset=offset)
                if self.transform is not None:
                    self.cov = self.transform(self.cov)

        if savestate:
            self.chain.detach_and_update(self.current)

        self.current['sample'].detach_()
        self.current['target_val'].detach_()
//  [python]
def predictW(gamma, Z, labels):
    Y = torch.zeros(n_runs,n_lsamples, n_ways,device='cuda')
    Y.scatter_(2,labels[:,:n_lsamples].unsqueeze(2), 1)
    delta = torch.sum(Z, 1)
    iden = torch.eye(5,device='cuda')
    iden = iden.reshape((1, 5, 5))
    iden = iden.repeat(10000, 1, 1)
    W = torch.bmm(torch.transpose(Z,1,2), Z/delta.unsqueeze(1))
    L = iden - W#(W + W.bmm(W))/2
    Z_l = Z[:,:n_lsamples]


    u = torch.cholesky(torch.bmm(torch.transpose(Z_l,1,2), Z_l) + gamma * L)# + 0.1*
    A = torch.cholesky_solve(torch.bmm(torch.transpose(Z_l,1,2), Y), u)
    P = Z.bmm(A)
    _, n, m = P.shape
    r = torch.ones(n_runs, n_lsamples + n_usamples,device='cuda')
    c = torch.ones(n_runs, n_ways,device='cuda') * (n_shot + n_queries)
    u = torch.zeros(n_runs, n).cuda()
    maxiters = 1000
    iters = 1
    # normalize this matrix
    while torch.max(torch.abs(u - P.sum(2))) > 0.01:
        u = P.sum(2)
        P *= (r / u).view((n_runs, -1, 1))
        P *= (c / P.sum(1)).view((n_runs, 1, -1))
        P[:,:n_lsamples].fill_(0)
        P[:,:n_lsamples].scatter_(2,labels[:,:n_lsamples].unsqueeze(2), 1)
        if iters == maxiters:
            break
        iters = iters + 1
    return P
//  [python]
def predictW(gamma, Z, labels):
    Y = torch.zeros(n_runs,n_lsamples, n_ways,device='cuda')
    Y.scatter_(2,labels[:,:n_lsamples].unsqueeze(2), 1)
    delta = torch.sum(Z, 1)
    iden = torch.eye(5,device='cuda')
    iden = iden.reshape((1, 5, 5))
    iden = iden.repeat(10000, 1, 1)
    W = torch.bmm(torch.transpose(Z,1,2), Z/delta.unsqueeze(1))
    L = iden - W#(W + W.bmm(W))/2
    Z_l = Z[:,:n_lsamples]


    u = torch.linalg.cholesky(torch.bmm(torch.transpose(Z_l,1,2), Z_l) + gamma * L)# + 0.1*
    A = torch.cholesky_solve(torch.bmm(torch.transpose(Z_l,1,2), Y), u)
    P = Z.bmm(A)
    _, n, m = P.shape
    r = torch.ones(n_runs, n_lsamples + n_usamples,device='cuda')
    c = torch.ones(n_runs, n_ways,device='cuda') * (n_shot + n_queries)
    u = torch.zeros(n_runs, n).cuda()
    maxiters = 1000
    iters = 1
    # normalize this matrix
    while torch.max(torch.abs(u - P.sum(2))) > 0.01:
        u = P.sum(2)
        P *= (r / u).view((n_runs, -1, 1))
        P *= (c / P.sum(1)).view((n_runs, 1, -1))
        P[:,:n_lsamples].fill_(0)
        P[:,:n_lsamples].scatter_(2,labels[:,:n_lsamples].unsqueeze(2), 1)
        if iters == maxiters:
            break
        iters = iters + 1
    return P
//  [python]
def jitchol(self, K, max_tries=5):
        # Cholesky decomposition + jitter // Inspired in the
        # Assure that is PSD
        try:
            L = torch.cholesky(K)
        except RuntimeError:
            diagK = torch.diag(K)
            if torch.any(diagK <= 0.0):
                raise ValueError("Non-positive diagonal elements in K")

            if K.shape[0] == K.shape[1]:
                jitter = torch.mean(diagK) * 1e-6
                num_tries = 1
                while num_tries <= max_tries and torch.isfinite(jitter):
                    try:
                        L = torch.cholesky(K + torch.eye(K.shape[0]) * jitter)
                        return L
                    except:
                        jitter *= 10
                    finally:
                        num_tries += 1
                raise ValueError("Non-positive definite, even with jitter")
        return L
//  [python]
def jitchol(self, K, max_tries=5):
        # Cholesky decomposition + jitter // Inspired in the
        # Assure that is PSD
        try:
            L = torch.linalg.cholesky(K)
        except RuntimeError:
            diagK = torch.diag(K)
            if torch.any(diagK <= 0.0):
                raise ValueError("Non-positive diagonal elements in K")

            if K.shape[0] == K.shape[1]:
                jitter = torch.mean(diagK) * 1e-6
                num_tries = 1
                while num_tries <= max_tries and torch.isfinite(jitter):
                    try:
                        L = torch.linalg.cholesky(K + torch.eye(K.shape[0]) * jitter)
                        return L
                    except:
                        jitter *= 10
                    finally:
                        num_tries += 1
                raise ValueError("Non-positive definite, even with jitter")
        return L
//  [python]
def cholesky(batched_matrix: Tensor):
    """Cholesky decomposition which attempts to regularize the
    matrix diagonal with some salt if the decomposition fails.

    Args:
        batched_matrix (Tensor): A 3-D tensor.

    Returns:
        Tensor: The cholesky factor.
    """
    try:
        cholesky_factor = torch.cholesky(batched_matrix)
    except RuntimeError:
        regularized_matrix = make_positive_definite(batched_matrix)
        cholesky_factor = torch.cholesky(regularized_matrix)
    return cholesky_factor
//  [python]
def cholesky(batched_matrix: Tensor):
    """Cholesky decomposition which attempts to regularize the
    matrix diagonal with some salt if the decomposition fails.

    Args:
        batched_matrix (Tensor): A 3-D tensor.

    Returns:
        Tensor: The cholesky factor.
    """
    try:
        cholesky_factor = torch.linalg.cholesky(batched_matrix)
    except RuntimeError:
        regularized_matrix = make_positive_definite(batched_matrix)
        cholesky_factor = torch.linalg.cholesky(regularized_matrix)
    return cholesky_factor
//  [python]
def diffusion(self, summary_type: str = "sample") -> Tensor:
        """Returns the diffusion matrix

        Args:
            summary_type (str, optional): "sample" or "mean"

        Returns:
            Tensor: the diffusion matrix
        """
        if summary_type == "mean":
            return torch.cholesky(self.diffusion_prior.process_noise.mean(0))
        elif summary_type == "sample":
            return torch.cholesky(self.diffusion_prior.process_noise)
        else:
            raise ValueError(f"summary_type {summary_type} not recognized.")
//  [python]
def diffusion(self, summary_type: str = "sample") -> Tensor:
        """Returns the diffusion matrix

        Args:
            summary_type (str, optional): "sample" or "mean"

        Returns:
            Tensor: the diffusion matrix
        """
        if summary_type == "mean":
            return torch.linalg.cholesky(self.diffusion_prior.process_noise.mean(0))
        elif summary_type == "sample":
            return torch.linalg.cholesky(self.diffusion_prior.process_noise)
        else:
            raise ValueError(f"summary_type {summary_type} not recognized.")
//  [python]
def log_det_by_cholesky(matrix: torch.Tensor) -> torch.Tensor:
    """Compute log det by cholesky.

    Args:
        matrix (torch.Tensor): matrix must be a positive define matrix.
        shape [N, C, D, D].
    Ref:
        https://github.com/tensorflow/tensorflow/blob/r1.13/tensorflow/python/ops/linalg/linalg_impl.py # pylint: disable=line-too-long
    """
    # This uses the property that the log det(A) = 2 * sum(log(real(diag(C))))
    # where C is the cholesky decomposition of A.
    chol = torch.cholesky(matrix)

    return 2.0 * torch.sum(
        torch.log(torch.diagonal(chol, dim1=-2, dim2=-1) + 1e-8), dim=-1
    )
//  [python]
def log_det_by_cholesky(matrix: torch.Tensor) -> torch.Tensor:
    """Compute log det by cholesky.

    Args:
        matrix (torch.Tensor): matrix must be a positive define matrix.
        shape [N, C, D, D].
    Ref:
        https://github.com/tensorflow/tensorflow/blob/r1.13/tensorflow/python/ops/linalg/linalg_impl.py # pylint: disable=line-too-long
    """
    # This uses the property that the log det(A) = 2 * sum(log(real(diag(C))))
    # where C is the cholesky decomposition of A.
    chol = torch.linalg.cholesky(matrix)

    return 2.0 * torch.sum(
        torch.log(torch.diagonal(chol, dim1=-2, dim2=-1) + 1e-8), dim=-1
    )
//  [python]
def solve_cholesky(A: th.Tensor, b: th.Tensor, w: th.Tensor) -> th.Tensor:
        # loss = |Ax - b|^2_w
        # Solve all Batch linear problems
        if A is None or b is None or w is None:
            return None

        # Check dimensionality, first is the batch dimension
        B, M = b.shape
        _, _, N = A.shape
        log.check_eq(th.Size([B, M, N]), A.shape)
        log.check_eq(th.Size([B, M]), w.shape) # We assume E diagonal: E=diag(w)
        WA = A * w.unsqueeze(-1) # multiply each row by the weight (uses broadcasting)
        AtWt = WA.transpose(1,2)
        AtEA = AtWt @ WA
        AtWtb = AtWt @ b.unsqueeze(-1)
        L = th.cholesky(AtEA)
        y = th.trtrs(AtWtb, L, upper=False)[0]
        x = th.trtrs(y, th.transpose(L, 1, 2), upper=True)[0]
        return x.squeeze(-1)
//  [python]
def solve_cholesky(A: th.Tensor, b: th.Tensor, w: th.Tensor) -> th.Tensor:
        # loss = |Ax - b|^2_w
        # Solve all Batch linear problems
        if A is None or b is None or w is None:
            return None

        # Check dimensionality, first is the batch dimension
        B, M = b.shape
        _, _, N = A.shape
        log.check_eq(th.Size([B, M, N]), A.shape)
        log.check_eq(th.Size([B, M]), w.shape) # We assume E diagonal: E=diag(w)
        WA = A * w.unsqueeze(-1) # multiply each row by the weight (uses broadcasting)
        AtWt = WA.transpose(1,2)
        AtEA = AtWt @ WA
        AtWtb = AtWt @ b.unsqueeze(-1)
        L = th.linalg.cholesky(AtEA)
        y = th.linalg.solve_triangular(L, AtWtb, upper=False)
        x = th.linalg.solve_triangular(th.transpose(L, 1, 2), y, upper=True)
        return x.squeeze(-1)
//  [python]
def solve_cholesky_11(A, b, E):
        # loss = |Ax - b|^2_sigma(\theta)
        At = th.transpose(A, 1, 2)
        AtE = At @ E
        AtEA = AtE @ A
        AtEb = AtE @ b
        L = th.cholesky(AtEA)
        x = th.cholesky_solve(AtEb, L, upper=False)
        return x
//  [python]
def solve_cholesky_11(A, b, E):
        # loss = |Ax - b|^2_sigma(\theta)
        At = th.transpose(A, 1, 2)
        AtE = At @ E
        AtEA = AtE @ A
        AtEb = AtE @ b
        L = th.linalg.cholesky(AtEA)
        x = th.linalg.cholesky_solve(AtEb, L, upper=False)
        return x
//  [python]
def is_psd(X):
    try:
        torch.cholesky(X)
    except RuntimeError as e:
        print(e)
        return False
    return True
//  [python]
def is_psd(X):
    try:
        torch.linalg.cholesky(X)
    except RuntimeError as e:
        print(e)
        return False
    return True
//  [python]
def make_psd(Kb,
             cholesky_tries=10,
             cholesky_perturb_init=1e-5,
             cholesky_perturb_scale=10):
    cholesky_perturb_factor = cholesky_perturb_init
    Kb_sqrt = None
    for ntry in range(cholesky_tries):
        try:
            Kbp = Kb + cholesky_perturb_factor * (
                torch.eye(Kb.shape[0], dtype=Kb.dtype, device=Kb.device) *
                torch.rand(Kb.shape[0], dtype=Kb.dtype, device=Kb.device)
                )
            Kb_sqrt = torch.cholesky(Kbp)
            break
        except RuntimeError as e:
            if ntry == cholesky_tries - 1:
                raise
            else:
                LOG.warning("Cholesky failed with perturb={} on error {}".format(cholesky_perturb_factor, str(e)))
                cholesky_perturb_factor = cholesky_perturb_factor * cholesky_perturb_scale
                continue

    return Kbp, Kb_sqrt
//  [python]
def make_psd(Kb,
             cholesky_tries=10,
             cholesky_perturb_init=1e-5,
             cholesky_perturb_scale=10):
    cholesky_perturb_factor = cholesky_perturb_init
    Kb_sqrt = None
    for ntry in range(cholesky_tries):
        try:
            Kbp = Kb + cholesky_perturb_factor * (
                torch.eye(Kb.shape[0], dtype=Kb.dtype, device=Kb.device) *
                torch.rand(Kb.shape[0], dtype=Kb.dtype, device=Kb.device)
                )
            Kb_sqrt = torch.linalg.cholesky(Kbp)
            break
        except RuntimeError as e:
            if ntry == cholesky_tries - 1:
                raise
            else:
                LOG.warning("Cholesky failed with perturb={} on error {}".format(cholesky_perturb_factor, str(e)))
                cholesky_perturb_factor = cholesky_perturb_factor * cholesky_perturb_scale
                continue

    return Kbp, Kb_sqrt