jiahuimbzuai/CoQuIR huggingface dataset:
//  [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)
    tZ_Z = torch.bmm(torch.transpose(Z,1,2), Z)
    delta = torch.sum(Z, 1)
    L = tZ_Z - torch.bmm(tZ_Z, tZ_Z/delta.unsqueeze(1))
    Z_l = Z[:,:n_lsamples]

    u = torch.cholesky(torch.bmm(torch.transpose(Z_l,1,2), Z_l) + gamma * L)
    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
    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):
    # #Certainty_scores = 1 + (Z*torch.log(Z)).sum(dim=2) / math.log(5)
    # Z[:,:n_lsamples].fill_(0)
    # Z[:,:n_lsamples].scatter_(2,labels[:,:n_lsamples].unsqueeze(2), 1)
    Y = torch.zeros(n_runs,n_lsamples, n_ways,device='cuda')
    Y.scatter_(2,labels[:,:n_lsamples].unsqueeze(2), 1)
    tZ_Z = torch.bmm(torch.transpose(Z,1,2), Z)
    delta = torch.sum(Z, 1)
    L = tZ_Z - torch.bmm(tZ_Z, tZ_Z/delta.unsqueeze(1))
    Z_l = Z[:,:n_lsamples]

    #A = np.dot(np.linalg.inv(torch.matmul(torch.transpose(Z_l,1,2), Z_l) + gamma * L), torch.bmm(torch.transpose(Z_l,1,2), Y))
    u = torch.linalg.cholesky(torch.bmm(torch.transpose(Z_l,1,2), Z_l) + gamma * L)# + 0.1*
    #u = torch.linalg.cholesky(gamma * L)
    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 _eval_kld_F(self, X):
        
        Kf = self.kerF.matrix(X)
        
        
        LN = torch.tril(self.chol_F_N)
        LK = torch.tril(self.chol_F_K)
        
        SN = torch.mm(LN, LN.T)
        SK = torch.mm(LK, LK.T)

        trace_term = torch.trace(torch.linalg.solve(Kf, SN))*torch.trace(SK)
        quad_term = torch.trace(torch.linalg.solve(Kf, torch.mm(self.Fmu, self.Fmu.T)))

        chol_Kf = torch.cholesky(Kf)
        logdet_term2 = self.dimK*torch.sum(torch.log(torch.square(torch.diag(chol_Kf))))
        logdet_term1 = self.dimK*torch.sum(torch.log(torch.square(torch.diag(LN)))) +\
                       self.dimN*torch.sum(torch.log(torch.square(torch.diag(LK))))
        
        
        kld_f = 0.5*(logdet_term2-logdet_term1-self.dimK*self.dimN+trace_term+quad_term)
        
        return kld_f
//  [python]
def _eval_kld_F(self, X):
        
        Kf = self.kerF.matrix(X)
        
        
        LN = torch.tril(self.chol_F_N)
        LK = torch.tril(self.chol_F_K)
        
        SN = torch.mm(LN, LN.T)
        SK = torch.mm(LK, LK.T)

        trace_term = torch.trace(torch.linalg.solve(Kf, SN))*torch.trace(SK)
        quad_term = torch.trace(torch.linalg.solve(Kf, torch.mm(self.Fmu, self.Fmu.T)))

        chol_Kf = torch.linalg.cholesky(Kf)
        logdet_term2 = self.dimK*torch.sum(torch.log(torch.square(torch.diag(chol_Kf))))
        logdet_term1 = self.dimK*torch.sum(torch.log(torch.square(torch.diag(LN)))) +\
                       self.dimN*torch.sum(torch.log(torch.square(torch.diag(LK))))
        
        
        kld_f = 0.5*(logdet_term2-logdet_term1-self.dimK*self.dimN+trace_term+quad_term)
        
        return kld_f
//  [python]
def M(self, q, eps=1e-6, verbose=False):
        """Returns a function that multiplies the mass matrix M(q) by a vector qdot
        Args:
            q: bs x D Tensor representing the position
        """
        assert q.ndim == 2
        lower_triangular = self.tril_Minv(q) * 0
        assert lower_triangular.ndim == 3

        def M_func(qdot):
            assert qdot.ndim == 2
            qdot = qdot.unsqueeze(-1)
            eps = 1
            diag_noise = eps * torch.eye(lower_triangular.size(-1), dtype=qdot.dtype, device=qdot.device)
            M = lower_triangular @ lower_triangular.transpose(-2, -1) + diag_noise
            M_cholesky = torch.cholesky(M)
            M_times_qdot = torch.cholesky_solve(qdot, M_cholesky).squeeze(-1)
            # M_times_qdot = torch.solve(qdot, M).solution.squeeze(-1)
            return M_times_qdot

        return M_func
//  [python]
def M(self, q, eps=1e-6, verbose=False):
        """Returns a function that multiplies the mass matrix M(q) by a vector qdot
        Args:
            q: bs x D Tensor representing the position
        """
        assert q.ndim == 2
        lower_triangular = self.tril_Minv(q) * 0
        assert lower_triangular.ndim == 3

        def M_func(qdot):
            assert qdot.ndim == 2
            qdot = qdot.unsqueeze(-1)
            eps = 1
            diag_noise = eps * torch.eye(lower_triangular.size(-1), dtype=qdot.dtype, device=qdot.device)
            M = lower_triangular @ lower_triangular.transpose(-2, -1) + diag_noise
            M_cholesky = torch.linalg.cholesky(M)
            M_times_qdot = torch.cholesky_solve(qdot, M_cholesky).squeeze(-1)
            # M_times_qdot = torch.solve(qdot, M).solution.squeeze(-1)
            return M_times_qdot

        return M_func
//  [python]
def _eval_kld_W(self, X):
        
        Kw = self.kerW.matrix(X)
        
         
        LN = torch.tril(self.chol_W_N)
        LK = torch.tril(self.chol_W_K)
        Ld1 = torch.tril(self.chol_W_d1)
        Ld2 = torch.tril(self.chol_W_d2)
        
        SN = torch.mm(LN, LN.T)
        SK = torch.mm(LK, LK.T)
        Sd1 = torch.mm(Ld1, Ld1.T)
        Sd2 = torch.mm(Ld2, Ld2.T)

         
        trace_term = torch.trace(torch.linalg.solve(Kw, SN))*torch.trace(SK)*torch.trace(Sd1)*torch.trace(Sd2)
        U = self.Wmu.flatten(1,-1)
        quad_term = torch.trace(torch.linalg.solve(Kw, torch.mm(U, U.T)))

        D = self.dimd1*self.dimd2*self.dimK*self.dimN
        chol_Kw = torch.cholesky(Kw)
        logdet_term2 = (D/self.dimN)*torch.sum(torch.log(torch.square(torch.diag(chol_Kw))))
        logdet_term1 = (D/self.dimN)*torch.sum(torch.log(torch.square(torch.diag(LN)))) +\
                       (D/self.dimK)*torch.sum(torch.log(torch.square(torch.diag(LK)))) +\
                       (D/self.dimd1)*torch.sum(torch.log(torch.square(torch.diag(Ld1)))) +\
                       (D/self.dimd2)*torch.sum(torch.log(torch.square(torch.diag(Ld2))))
        
        kld_w = 0.5*(logdet_term2-logdet_term1-D+trace_term+quad_term)
        
        return kld_w
//  [python]
def _eval_kld_W(self, X):
        
        Kw = self.kerW.matrix(X)
        
         
        LN = torch.tril(self.chol_W_N)
        LK = torch.tril(self.chol_W_K)
        Ld1 = torch.tril(self.chol_W_d1)
        Ld2 = torch.tril(self.chol_W_d2)
        
        SN = torch.mm(LN, LN.T)
        SK = torch.mm(LK, LK.T)
        Sd1 = torch.mm(Ld1, Ld1.T)
        Sd2 = torch.mm(Ld2, Ld2.T)

         
        trace_term = torch.trace(torch.linalg.solve(Kw, SN))*torch.trace(SK)*torch.trace(Sd1)*torch.trace(Sd2)
        U = self.Wmu.flatten(1,-1)
        quad_term = torch.trace(torch.linalg.solve(Kw, torch.mm(U, U.T)))

        D = self.dimd1*self.dimd2*self.dimK*self.dimN
        chol_Kw = torch.linalg.cholesky(Kw)
        logdet_term2 = (D/self.dimN)*torch.sum(torch.log(torch.square(torch.diag(chol_Kw))))
        logdet_term1 = (D/self.dimN)*torch.sum(torch.log(torch.square(torch.diag(LN)))) +\
                       (D/self.dimK)*torch.sum(torch.log(torch.square(torch.diag(LK)))) +\
                       (D/self.dimd1)*torch.sum(torch.log(torch.square(torch.diag(Ld1)))) +\
                       (D/self.dimd2)*torch.sum(torch.log(torch.square(torch.diag(Ld2))))
        
        kld_w = 0.5*(logdet_term2-logdet_term1-D+trace_term+quad_term)
        
        return kld_w
//  [python]
def fit_normal(self):
        """
        Fit a multivariate normal to the live points
        Returns
        -------
        mvn : torch.distributions.MultivariateNormal
            The multivariate normal distribution
        """
        x = self.live_points.get_values()
        mean = torch.mean(x, dim=0)
        cov = (1/self.eff)*torch.cov(x.T)
        assert torch.linalg.det(cov) > 0, "Covariance not positive semidefinite"
        mvn = torch.distributions.MultivariateNormal(mean,
                                                     scale_tril=torch.cholesky(cov).to(self.device))
        return mvn
//  [python]
def fit_normal(self):
        """
        Fit a multivariate normal to the live points
        Returns
        -------
        mvn : torch.distributions.MultivariateNormal
            The multivariate normal distribution
        """
        x = self.live_points.get_values()
        mean = torch.mean(x, dim=0)
        cov = (1/self.eff)*torch.cov(x.T)
        assert torch.linalg.det(cov) > 0, "Covariance not positive semidefinite"
        mvn = torch.distributions.MultivariateNormal(mean,
                                                     scale_tril=torch.linalg.cholesky(cov).to(self.device))
        return mvn
//  [python]
def cholesky_inverse(fish, momentum):
    """Performs the inverse of a matrix, using the cholesky inverse (with the vector).

    Parameters
    ----------
    fish : torch.tensor
        Square matrix to be inverted: shape (D,D).
    momentum : torch.tensor
        Vector of shape (D,).

    Returns
    -------
    torch.tensor
        Returns the inverted matrix multiplied by the vector.

    """
    lower = torch.cholesky(fish)
    y = torch.triangular_solve(momentum.view(-1, 1), lower, upper=False, unitriangular=False)[0]
    fish_inv_p = torch.triangular_solve(y, lower.t(), upper=True, unitriangular=False)[0]
    return fish_inv_p
//  [python]
def cholesky_inverse(fish, momentum):
    """Performs the inverse of a matrix, using the cholesky inverse (with the vector).

    Parameters
    ----------
    fish : torch.tensor
        Square matrix to be inverted: shape (D,D).
    momentum : torch.tensor
        Vector of shape (D,).

    Returns
    -------
    torch.tensor
        Returns the inverted matrix multiplied by the vector.

    """
    lower = torch.linalg.cholesky(fish)
    y = torch.linalg.solve_triangular(lower, momentum.view(-1, 1), upper=False, unitriangular=False)
    fish_inv_p = torch.linalg.solve_triangular(lower.t(), y, upper=True, unitriangular=False)
    return fish_inv_p
//  [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)
        proposed['sample'] = self.current['sample'].clone().detach() + self.chol_cov @ 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
        else:
            self.model.set_params(self.current['sample'].clone().detach())
            self.current['accepted'] = 0

        h = min(1, self.model.num_params() * (self.counter.idx + 1 - offset) ** (-self.g))
        self.chol_cov = torch.cholesky(self.chol_cov @ (
            torch.eye(self.model.num_params(), dtype=self.model.dtype, device=self.model.device) + \
            h * (min(1, torch.exp(log_rate).item()) - self.a
            ) * torch.ger(randn_sample, randn_sample) / \
            torch.dot(randn_sample, randn_sample).item()) @ self.chol_cov.t())

        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)
        proposed['sample'] = self.current['sample'].clone().detach() + self.chol_cov @ 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
        else:
            self.model.set_params(self.current['sample'].clone().detach())
            self.current['accepted'] = 0

        h = min(1, self.model.num_params() * (self.counter.idx + 1 - offset) ** (-self.g))
        self.chol_cov = torch.linalg.cholesky(self.chol_cov @ (
            torch.eye(self.model.num_params(), dtype=self.model.dtype, device=self.model.device) + \
            h * (min(1, torch.exp(log_rate).item()) - self.a
            ) * torch.ger(randn_sample, randn_sample) / \
            torch.dot(randn_sample, randn_sample).item()) @ self.chol_cov.t())

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

        self.current['sample'].detach_()
        self.current['target_val'].detach_()
//  [python]
def _inverse(self, y):
        y = t.cholesky(y)
        return LowerCholeskyTransform().inv(y)
//  [python]
def _inverse(self, y):
        y = torch.linalg.cholesky(y)
        return LowerCholeskyTransform().inv(y)
//  [python]
def get_objective(self, Y, Z, aux_data, **kwargs):
        # TODO: look at either torch.bmm or torch.matmul
        covar_mat = aux_data
        covar_mat_Z_t = (torch.cholesky(covar_mat) * Z.unsqueeze(dim=-2)).sum(dim=-1)
        quad_term = covar_mat_Z_t.square().sum(dim=-1)
        obj = (Y * Z).sum(dim=-1) - self.alpha * quad_term
        return obj
//  [python]
def get_objective(self, Y, Z, aux_data, **kwargs):
        # TODO: look at either torch.bmm or torch.matmul
        covar_mat = aux_data
        covar_mat_Z_t = (torch.linalg.cholesky(covar_mat) * Z.unsqueeze(dim=-2)).sum(dim=-1)
        quad_term = covar_mat_Z_t.square().sum(dim=-1)
        obj = (Y * Z).sum(dim=-1) - self.alpha * quad_term
        return obj
//  [python]
def get_decision(self, Y, aux_data, max_instances_per_batch=1500, **kwargs):
        # Get the sqrt of the covariance matrix
        covar_mat = aux_data
        sqrt_covar = torch.cholesky(covar_mat)

        # Split Y into reasonably sized chunks so that we don't run into memory issues
        # Assumption Y is only 2D at max
        assert Y.ndim <= 2
        if Y.ndim == 2:
            results = []
            for start in range(0, Y.shape[0], max_instances_per_batch):
                end = min(Y.shape[0], start + max_instances_per_batch)
                result = self.opt(Y[start:end], sqrt_covar)[0]
                results.append(result)
            return torch.cat(results, dim=0)
        else:
            return self.opt(Y, sqrt_covar)[0]
//  [python]
def get_decision(self, Y, aux_data, max_instances_per_batch=1500, **kwargs):
        # Get the sqrt of the covariance matrix
        covar_mat = aux_data
        sqrt_covar = torch.linalg.cholesky(covar_mat)

        # Split Y into reasonably sized chunks so that we don't run into memory issues
        # Assumption Y is only 2D at max
        assert Y.ndim <= 2
        if Y.ndim == 2:
            results = []
            for start in range(0, Y.shape[0], max_instances_per_batch):
                end = min(Y.shape[0], start + max_instances_per_batch)
                result = self.opt(Y[start:end], sqrt_covar)[0]
                results.append(result)
            return torch.cat(results, dim=0)
        else:
            return self.opt(Y, sqrt_covar)[0]
//  [python]
def forward(
        self,
        x: TensorType[..., 'dim'],
        i: TensorType[..., 1],
        **kwargs,
    ) -> Tuple[TensorType[..., 'dim'], TensorType[..., 'dim']]:

        noise_gaussian = torch.randn_like(x)

        if self.is_time_series:
            cov = self.noise.covariance(**kwargs)
            L = torch.cholesky(cov)
            noise = L @ noise_gaussian
        else:
            noise = noise_gaussian

        alpha = self.alphas[i.long()].to(x)
        y = torch.sqrt(alpha) * x + torch.sqrt(1 - alpha) * noise

        if self.predict_gaussian_noise:
            return y, noise_gaussian
        else:
            return y, noise
//  [python]
def forward(
        self,
        x: TensorType[..., 'dim'],
        i: TensorType[..., 1],
        **kwargs,
    ) -> Tuple[TensorType[..., 'dim'], TensorType[..., 'dim']]:

        noise_gaussian = torch.randn_like(x)

        if self.is_time_series:
            cov = self.noise.covariance(**kwargs)
            L = torch.linalg.cholesky(cov)
            noise = L @ noise_gaussian
        else:
            noise = noise_gaussian

        alpha = self.alphas[i.long()].to(x)
        y = torch.sqrt(alpha) * x + torch.sqrt(1 - alpha) * noise

        if self.predict_gaussian_noise:
            return y, noise_gaussian
        else:
            return y, noise
//  [python]
def fasterquant(
        self, blocksize=128, percdamp=.01, groupsize=-1, actorder=False
    ):
        W = self.weight.clone()
        W = W.float()
        tick = time.time()

        if not self.quantizer.ready():
            self.quantizer.find_params(W)

        H = self.H
        del self.H
        dead = torch.diag(H) == 0
        H[dead, dead] = 1

        if actorder:
            perm = torch.argsort(torch.diag(H), descending=True)
            W = W[perm, :]
            H = H[perm][:, perm]

        damp = percdamp * torch.mean(torch.diag(H))
        diag = torch.arange(self.rows)
        H[diag, diag] += damp
        H = torch.cholesky(H)
        H = torch.cholesky_inverse(H)
        H = torch.cholesky(H)
        Hinv = H

        Losses = torch.zeros_like(W)
        Q = torch.zeros_like(W)
        Qint = torch.zeros_like(W)

        for i1 in range(0, self.rows, blocksize):
            i2 = min(i1 + blocksize, self.rows)
            count = i2 - i1

            W1 = W[i1:i2, :].clone()
            Err1 = torch.zeros_like(W1)

            for i in range(count):
                w = W1[i, :]
                d = Hinv[i1 + i, i1 + i]

                if groupsize != -1:
                    if (i1 + i) % groupsize == 0:
                        self.quantizer.find_params(W[(i1 + i):(i1 + i + groupsize), :])

                q_int = quantize_to_int(w, self.quantizer.scale, self.quantizer.zero, self.quantizer.maxq)
                Qint[i + i1, :] = q_int
                q = dequantize_to_float(q_int, self.quantizer.scale, self.quantizer.zero)
                Q[i + i1, :] = q

                Losses[i + i1, :] = (w - q) ** 2 / ( 2 * d ** 2)
                err1 = (w - q) / d
                Err1[i, :] = err1

                h = Hinv[(i1 + i):i2, i1 + i]
                err1 = err1.reshape([1, err1.shape[0]])
                h = h.reshape([h.shape[0], 1])
                W1[i:, :] -= h.matmul(err1)

            W[i2:, :] -= Hinv[i2:, i1:i2].matmul(Err1)

        if actorder:
            invperm = torch.argsort(perm)
            Q = Q[invperm, :]

        # if isinstance(self.layer, transformers.Conv1D):
        #     Q = Q.t()

        print('time %.2f' % (time.time() - tick))
        print('error', torch.sum(Losses).item())

        print("Input Weight (float):", self.weight)
        print("Output quantized weight (int{})".format(self.wbits), Qint)
        print("Output quantized weight (float)", Q)
        return Qint, self.quantizer.scale, self.quantizer.zero
//  [python]
def fasterquant(
        self, blocksize=128, percdamp=.01, groupsize=-1, actorder=False
    ):
        W = self.weight.clone()
        W = W.float()
        tick = time.time()

        if not self.quantizer.ready():
            self.quantizer.find_params(W)

        H = self.H
        del self.H
        dead = torch.diag(H) == 0
        H[dead, dead] = 1

        if actorder:
            perm = torch.argsort(torch.diag(H), descending=True)
            W = W[perm, :]
            H = H[perm][:, perm]

        damp = percdamp * torch.mean(torch.diag(H))
        diag = torch.arange(self.rows)
        H[diag, diag] += damp
        H = torch.linalg.cholesky(H)
        H = torch.cholesky_inverse(H)
        H = torch.linalg.cholesky(H)
        Hinv = H

        Losses = torch.zeros_like(W)
        Q = torch.zeros_like(W)
        Qint = torch.zeros_like(W)

        for i1 in range(0, self.rows, blocksize):
            i2 = min(i1 + blocksize, self.rows)
            count = i2 - i1

            W1 = W[i1:i2, :].clone()
            Err1 = torch.zeros_like(W1)

            for i in range(count):
                w = W1[i, :]
                d = Hinv[i1 + i, i1 + i]

                if groupsize != -1:
                    if (i1 + i) % groupsize == 0:
                        self.quantizer.find_params(W[(i1 + i):(i1 + i + groupsize), :])

                q_int = quantize_to_int(w, self.quantizer.scale, self.quantizer.zero, self.quantizer.maxq)
                Qint[i + i1, :] = q_int
                q = dequantize_to_float(q_int, self.quantizer.scale, self.quantizer.zero)
                Q[i + i1, :] = q

                Losses[i + i1, :] = (w - q) ** 2 / ( 2 * d ** 2)
                err1 = (w - q) / d
                Err1[i, :] = err1

                h = Hinv[(i1 + i):i2, i1 + i]
                err1 = err1.reshape([1, err1.shape[0]])
                h = h.reshape([h.shape[0], 1])
                W1[i:, :] -= h.matmul(err1)

            W[i2:, :] -= Hinv[i2:, i1:i2].matmul(Err1)

        if actorder:
            invperm = torch.argsort(perm)
            Q = Q[invperm, :]

        # if isinstance(self.layer, transformers.Conv1D):
        #     Q = Q.t()

        print('time %.2f' % (time.time() - tick))
        print('error', torch.sum(Losses).item())

        print("Input Weight (float):", self.weight)
        print("Output quantized weight (int{})".format(self.wbits), Qint)
        print("Output quantized weight (float)", Q)
        return Qint, self.quantizer.scale, self.quantizer.zero
//  [python]
def test_psd_safe_cholesky_pd(self, cuda=False):
        device = torch.device("cuda") if cuda else torch.device("cpu")
        for dtype in (torch.float, torch.double):
            for batch_mode in (False, True):
                if batch_mode:
                    A = self._gen_test_psd().to(device=device, dtype=dtype)
                    D = torch.eye(2).type_as(A).unsqueeze(0).repeat(2, 1, 1)
                else:
                    A = self._gen_test_psd()[0].to(device=device, dtype=dtype)
                    D = torch.eye(2).type_as(A)
                A += D
                # basic
                L = torch.cholesky(A)
                L_safe = psd_safe_cholesky(A)
                self.assertTrue(torch.allclose(L, L_safe))
                # upper
                L = torch.cholesky(A).mT
                L_safe = psd_safe_cholesky(A, upper=True)
                self.assertTrue(torch.allclose(L, L_safe))
                # output tensors
                L = torch.empty_like(A)
                L_safe = torch.empty_like(A)
                torch.cholesky(A, out=L)
                psd_safe_cholesky(A, out=L_safe)
                self.assertTrue(torch.allclose(L, L_safe))
                # output tensors, upper
                torch.cholesky(A, out=L)
                L = L.mT
                psd_safe_cholesky(A, upper=True, out=L_safe)
                self.assertTrue(torch.allclose(L, L_safe))
                # make sure jitter doesn't do anything if p.d.
                L = torch.cholesky(A)
                L_safe = psd_safe_cholesky(A, jitter=1e-2)
                self.assertTrue(torch.allclose(L, L_safe))
//  [python]
def test_psd_safe_cholesky_pd(self, cuda=False):
        device = torch.device("cuda") if cuda else torch.device("cpu")
        for dtype in (torch.float, torch.double):
            for batch_mode in (False, True):
                if batch_mode:
                    A = self._gen_test_psd().to(device=device, dtype=dtype)
                    D = torch.eye(2).type_as(A).unsqueeze(0).repeat(2, 1, 1)
                else:
                    A = self._gen_test_psd()[0].to(device=device, dtype=dtype)
                    D = torch.eye(2).type_as(A)
                A += D
                # basic
                L = torch.linalg.cholesky(A)
                L_safe = psd_safe_cholesky(A)
                self.assertTrue(torch.allclose(L, L_safe))
                # upper
                L = torch.linalg.cholesky(A).mT
                L_safe = psd_safe_cholesky(A, upper=True)
                self.assertTrue(torch.allclose(L, L_safe))
                # output tensors
                L = torch.empty_like(A)
                L_safe = torch.empty_like(A)
                torch.linalg.cholesky(A, out=L)
                psd_safe_cholesky(A, out=L_safe)
                self.assertTrue(torch.allclose(L, L_safe))
                # output tensors, upper
                torch.linalg.cholesky(A, out=L)
                L = L.mT
                psd_safe_cholesky(A, upper=True, out=L_safe)
                self.assertTrue(torch.allclose(L, L_safe))
                # make sure jitter doesn't do anything if p.d.
                L = torch.linalg.cholesky(A)
                L_safe = psd_safe_cholesky(A, jitter=1e-2)
                self.assertTrue(torch.allclose(L, L_safe))
//  [python]
def log_prob(
        self,
        model: Callable,
        x: TensorType[..., 'dim'],
        num_samples: int = 1,
        **kwargs,
    ) -> TensorType[..., 1]:
        if self.is_time_series and self.predict_gaussian_noise:
            cov = self.noise.covariance(**kwargs)
            L = torch.cholesky(cov)
        else:
            L = None

        func = self._elbo_parallel if self.parallel_elbo else self._elbo_sequential
        return func(model, x, num_samples=num_samples, L=L, **kwargs)
//  [python]
def log_prob(
        self,
        model: Callable,
        x: TensorType[..., 'dim'],
        num_samples: int = 1,
        **kwargs,
    ) -> TensorType[..., 1]:
        if self.is_time_series and self.predict_gaussian_noise:
            cov = self.noise.covariance(**kwargs)
            L = torch.linalg.cholesky(cov)
        else:
            L = None

        func = self._elbo_parallel if self.parallel_elbo else self._elbo_sequential
        return func(model, x, num_samples=num_samples, L=L, **kwargs)
//  [python]
def sample(
        self,
        model: Callable,
        num_samples: Union[int, Tuple],
        device: str = 'cpu',
        **kwargs,
    ) -> TensorType['*num_samples', 'dim']:
        if isinstance(num_samples, int):
            num_samples = (num_samples,)

        x = self.noise(*num_samples, **kwargs).to(device)

        if self.is_time_series and self.predict_gaussian_noise:
            cov = self.noise.covariance(**kwargs)
            L = torch.cholesky(cov)
        else:
            L = None

        for diff_step in reversed(range(0, self.num_steps)):
            alpha = self.alphas[diff_step]
            beta = self.betas[diff_step]

            # An alternative can be:
            # alpha_prev = self.alphas[diff_step - 1]
            # sigma = beta * (1 - alpha_prev) / (1 - alpha)
            sigma = beta

            if diff_step == 0:
                z = 0
            else:
                z = self.noise(*num_samples, **kwargs).to(device)

            i = torch.Tensor([diff_step]).expand_as(x[...,:1]).to(device)
            pred_noise = model(x, i=i, **kwargs)

            if L is not None:
                pred_noise = L @ pred_noise

            x = (x - beta * pred_noise / (1 - alpha).sqrt()) / (1 - beta).sqrt() + sigma.sqrt() * z

        return x
//  [python]
def sample(
        self,
        model: Callable,
        num_samples: Union[int, Tuple],
        device: str = 'cpu',
        **kwargs,
    ) -> TensorType['*num_samples', 'dim']:
        if isinstance(num_samples, int):
            num_samples = (num_samples,)

        x = self.noise(*num_samples, **kwargs).to(device)

        if self.is_time_series and self.predict_gaussian_noise:
            cov = self.noise.covariance(**kwargs)
            L = torch.linalg.cholesky(cov)
        else:
            L = None

        for diff_step in reversed(range(0, self.num_steps)):
            alpha = self.alphas[diff_step]
            beta = self.betas[diff_step]

            # An alternative can be:
            # alpha_prev = self.alphas[diff_step - 1]
            # sigma = beta * (1 - alpha_prev) / (1 - alpha)
            sigma = beta

            if diff_step == 0:
                z = 0
            else:
                z = self.noise(*num_samples, **kwargs).to(device)

            i = torch.Tensor([diff_step]).expand_as(x[...,:1]).to(device)
            pred_noise = model(x, i=i, **kwargs)

            if L is not None:
                pred_noise = L @ pred_noise

            x = (x - beta * pred_noise / (1 - alpha).sqrt()) / (1 - beta).sqrt() + sigma.sqrt() * z

        return x
//  [python]
def test_psd_safe_cholesky_psd(self, cuda=False):
        device = torch.device("cuda") if cuda else torch.device("cpu")
        for dtype in (torch.float, torch.double):
            for batch_mode in (False, True):
                if batch_mode:
                    A = self._gen_test_psd().to(device=device, dtype=dtype)
                else:
                    A = self._gen_test_psd()[0].to(device=device, dtype=dtype)
                idx = torch.arange(A.shape[-1], device=A.device)
                # default values
                Aprime = A.clone()
                Aprime[..., idx, idx] += 1e-6 if A.dtype == torch.float32 else 1e-8
                L_exp = torch.cholesky(Aprime)
                with warnings.catch_warnings(record=True) as w:
                    # Makes sure warnings we catch don't cause `-w error` to fail
                    warnings.simplefilter("always", NumericalWarning)

                    L_safe = psd_safe_cholesky(A)
                    self.assertTrue(any(issubclass(w_.category, NumericalWarning) for w_ in w))
                    self.assertTrue(any("A not p.d., added jitter" in str(w_.message) for w_ in w))
                self.assertTrue(torch.allclose(L_exp, L_safe))
                # user-defined value
                Aprime = A.clone()
                Aprime[..., idx, idx] += 1e-2
                L_exp = torch.cholesky(Aprime)
                with warnings.catch_warnings(record=True) as w:
                    # Makes sure warnings we catch don't cause `-w error` to fail
                    warnings.simplefilter("always", NumericalWarning)

                    L_safe = psd_safe_cholesky(A, jitter=1e-2)
                    self.assertTrue(any(issubclass(w_.category, NumericalWarning) for w_ in w))
                    self.assertTrue(any("A not p.d., added jitter" in str(w_.message) for w_ in w))
                self.assertTrue(torch.allclose(L_exp, L_safe))
//  [python]
def test_psd_safe_cholesky_psd(self, cuda=False):
        device = torch.device("cuda") if cuda else torch.device("cpu")
        for dtype in (torch.float, torch.double):
            for batch_mode in (False, True):
                if batch_mode:
                    A = self._gen_test_psd().to(device=device, dtype=dtype)
                else:
                    A = self._gen_test_psd()[0].to(device=device, dtype=dtype)
                idx = torch.arange(A.shape[-1], device=A.device)
                # default values
                Aprime = A.clone()
                Aprime[..., idx, idx] += 1e-6 if A.dtype == torch.float32 else 1e-8
                L_exp = torch.linalg.cholesky(Aprime)
                with warnings.catch_warnings(record=True) as w:
                    # Makes sure warnings we catch don't cause `-w error` to fail
                    warnings.simplefilter("always", NumericalWarning)

                    L_safe = psd_safe_cholesky(A)
                    self.assertTrue(any(issubclass(w_.category, NumericalWarning) for w_ in w))
                    self.assertTrue(any("A not p.d., added jitter" in str(w_.message) for w_ in w))
                self.assertTrue(torch.allclose(L_exp, L_safe))
                # user-defined value
                Aprime = A.clone()
                Aprime[..., idx, idx] += 1e-2
                L_exp = torch.linalg.cholesky(Aprime)
                with warnings.catch_warnings(record=True) as w:
                    # Makes sure warnings we catch don't cause `-w error` to fail
                    warnings.simplefilter("always", NumericalWarning)

                    L_safe = psd_safe_cholesky(A, jitter=1e-2)
                    self.assertTrue(any(issubclass(w_.category, NumericalWarning) for w_ in w))
                    self.assertTrue(any("A not p.d., added jitter" in str(w_.message) for w_ in w))
                self.assertTrue(torch.allclose(L_exp, L_safe))
//  [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):
            try:
                log_det[k] = 2 * torch.log(torch.diagonal(torch.cholesky(var[0,k]))).sum()
            except:
                print("cholesky needs positive semi-definite, adding 1e-3 now")
                log_det[k] = 2 * torch.log(torch.diagonal(torch.cholesky(var[0, k] + torch.eye(len(var[0, k]), device=var.device) * 1e-3))).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):
            try:
                log_det[k] = 2 * torch.log(torch.diagonal(torch.linalg.cholesky(var[0,k]))).sum()
            except:
                print("cholesky needs positive semi-definite, adding 1e-3 now")
                log_det[k] = 2 * torch.log(torch.diagonal(torch.linalg.cholesky(var[0, k] + torch.eye(len(var[0, k]), device=var.device) * 1e-3))).sum()

        return log_det.unsqueeze(-1)
//  [python]
def is_psd(self, mat):
        """
        Args:
           - mat: torch.tensor, symmetric matrix

        Returns:
           - flag: bool, flag to judge whether or not the given matrix is positive semi-definite
        """
        try:
            torch.cholesky(mat)
            return bool((mat == mat.T).all() and (torch.linalg.eig(mat)[0].real >= 0).all())
        except:
            return False
//  [python]
def is_psd(self, mat):
        """
        Args:
           - mat: torch.tensor, symmetric matrix

        Returns:
           - flag: bool, flag to judge whether or not the given matrix is positive semi-definite
        """
        try:
            torch.linalg.cholesky(mat)
            return bool((mat == mat.T).all() and (torch.linalg.eig(mat)[0].real >= 0).all())
        except:
            return False
//  [python]
def negative_log_likelihood(self):
        y_num, y_dimension = self.y.shape
        sigma = self.kernel_customized(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, 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_customized(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, 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 __init__(self, s):
        n, m = s.shape
        if n != m:
            raise ValueError("Covariance matrix is not symmetric.")
        self.n = n
        self.chol = torch.cholesky(s, upper=False)
//  [python]
def __init__(self, s):
        n, m = s.shape
        if n != m:
            raise ValueError("Covariance matrix is not symmetric.")
        self.n = n
        self.chol = torch.linalg.cholesky(s, lower=True)
//  [python]
def model(self):
        self.set_mode("model")

        M = self.Xu.size(0)
        Kuu = self.kernel(self.Xu).contiguous()
        Kuu.view(-1)[:: M + 1] += self.jitter  # add jitter to the diagonal
        Luu = torch.cholesky(Kuu)

        zero_loc = self.Xu.new_zeros(self.u_loc.shape)
        if self.whiten:
            identity = eye_like(self.Xu, M)
            pyro.sample(
                self._pyro_get_fullname("u"),
                dist.MultivariateNormal(zero_loc, scale_tril=identity).to_event(
                    zero_loc.dim() - 1
                ),
            )
        else:
            pyro.sample(
                self._pyro_get_fullname("u"),
                dist.MultivariateNormal(zero_loc, scale_tril=Luu).to_event(
                    zero_loc.dim() - 1
                ),
            )

        f_loc, f_var = conditional(
            self.X,
            self.Xu,
            self.kernel,
            self.u_loc,
            self.u_scale_tril,
            Luu,
            full_cov=False,
            whiten=self.whiten,
            jitter=self.jitter,
        )

        f_loc = f_loc + self.mean_function(self.X)
        if self.y is None:
            return f_loc, f_var
        else:
            # we would like to load likelihood's parameters outside poutine.scale context
            self.likelihood._load_pyro_samples()
            with poutine.scale(scale=self.num_data / self.X.size(0)):
                return self.likelihood(f_loc, f_var, self.y)
//  [python]
def model(self):
        self.set_mode("model")

        M = self.Xu.size(0)
        Kuu = self.kernel(self.Xu).contiguous()
        Kuu.view(-1)[:: M + 1] += self.jitter  # add jitter to the diagonal
        Luu = torch.linalg.cholesky(Kuu)

        zero_loc = self.Xu.new_zeros(self.u_loc.shape)
        if self.whiten:
            identity = eye_like(self.Xu, M)
            pyro.sample(
                self._pyro_get_fullname("u"),
                dist.MultivariateNormal(zero_loc, scale_tril=identity).to_event(
                    zero_loc.dim() - 1
                ),
            )
        else:
            pyro.sample(
                self._pyro_get_fullname("u"),
                dist.MultivariateNormal(zero_loc, scale_tril=Luu).to_event(
                    zero_loc.dim() - 1
                ),
            )

        f_loc, f_var = conditional(
            self.X,
            self.Xu,
            self.kernel,
            self.u_loc,
            self.u_scale_tril,
            Luu,
            full_cov=False,
            whiten=self.whiten,
            jitter=self.jitter,
        )

        f_loc = f_loc + self.mean_function(self.X)
        if self.y is None:
            return f_loc, f_var
        else:
            # we would like to load likelihood's parameters outside poutine.scale context
            self.likelihood._load_pyro_samples()
            with poutine.scale(scale=self.num_data / self.X.size(0)):
                return self.likelihood(f_loc, f_var, self.y)
//  [python]
def negative_log_likelihood(kernel, log_beta, x_train, y_train):
        
    ## A scheme proposed by WJ to make GP more stable -2024/3/4

    y_num, y_dimension = y_train.shape
    Sigma = kernel(x_train, x_train) + log_beta.exp().pow(-1) * torch.eye(
        x_train.size(0)) + JITTER * kernel(x_train, x_train).mean() * torch.eye(x_train.size(0))
    
    L = torch.cholesky(Sigma)
    #option 1 (use this if torch supports)
    Gamma,_ = torch.triangular_solve(y_train, 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(kernel, log_beta, x_train, y_train):
        
    ## A scheme proposed by WJ to make GP more stable -2024/3/4

    y_num, y_dimension = y_train.shape
    Sigma = kernel(x_train, x_train) + log_beta.exp().pow(-1) * torch.eye(
        x_train.size(0)) + JITTER * kernel(x_train, x_train).mean() * torch.eye(x_train.size(0))
    
    L = torch.linalg.cholesky(Sigma)
    #option 1 (use this if torch supports)
    Gamma,_ = torch.triangular_solve(y_train, 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 conditional_Gaussian(y, Sigma, K_s, K_ss, Kinv_method='cholesky3'):
    # Sigma = Sigma + torch.eye(len(Sigma)) * EPS
    if Kinv_method == 'cholesky1':   # kernel inverse is not stable, use cholesky decomposition instead
        L = torch.cholesky(Sigma)
        L_inv = torch.inverse(L)
        K_inv = L_inv.T @ L_inv
        alpha = K_inv @ y
        mu = K_s.T @ alpha
        v = L_inv @ K_s
        cov = K_ss - v.T @ v
    elif Kinv_method == 'cholesky3':
        # recommended implementation, fastest so far
        L = torch.cholesky(Sigma)
        alpha = torch.cholesky_solve(y, L)
        mu = K_s.T @ alpha
        # v = torch.cholesky_solve(K_s, L)    # wrong implementation
        v = L.inverse() @ K_s   # correct implementation
        cov = K_ss - v.T @ v
    elif Kinv_method == 'direct':
        K_inv = torch.inverse(Sigma)
        mu = K_s.T @ K_inv @ y
        cov = K_ss - K_s.T @ K_inv @ K_s
    else:
        raise ValueError('Kinv_method should be either direct or cholesky')
    
    return mu, cov
//  [python]
def conditional_Gaussian(y, Sigma, K_s, K_ss, Kinv_method='cholesky3'):
    # Sigma = Sigma + torch.eye(len(Sigma)) * EPS
    if Kinv_method == 'cholesky1':   # kernel inverse is not stable, use cholesky decomposition instead
        L = torch.linalg.cholesky(Sigma)
        L_inv = torch.inverse(L)
        K_inv = L_inv.T @ L_inv
        alpha = K_inv @ y
        mu = K_s.T @ alpha
        v = L_inv @ K_s
        cov = K_ss - v.T @ v
    elif Kinv_method == 'cholesky3':
        # recommended implementation, fastest so far
        L = torch.linalg.cholesky(Sigma)
        alpha = torch.cholesky_solve(y, L)
        mu = K_s.T @ alpha
        # v = torch.cholesky_solve(K_s, L)    # wrong implementation
        v = L.inverse() @ K_s   # correct implementation
        cov = K_ss - v.T @ v
    elif Kinv_method == 'direct':
        K_inv = torch.inverse(Sigma)
        mu = K_s.T @ K_inv @ y
        cov = K_ss - K_s.T @ K_inv @ K_s
    else:
        raise ValueError('Kinv_method should be either direct or cholesky')
    
    return mu, cov
//  [python]
def test_pyro_catch_error(self):
        def potential_fn(z):
            mvn = pyro.distributions.MultivariateNormal(
                loc=torch.zeros(2),
                covariance_matrix=z["K"],
            )
            return mvn.log_prob(torch.zeros(2))

        # Test base case where everything is fine
        z = {"K": torch.eye(2)}
        grads, val = potential_grad(potential_fn, z)
        self.assertAllClose(grads["K"], -0.5 * torch.eye(2))
        norm_mvn = torch.distributions.Normal(0, 1)
        self.assertAllClose(val, 2 * norm_mvn.log_prob(torch.tensor(0.0)))

        # Default behavior should catch the ValueError when trying to instantiate
        # the MVN and return NaN instead
        z = {"K": torch.ones(2, 2)}
        _, val = potential_grad(potential_fn, z)
        self.assertTrue(torch.isnan(val))

        # Default behavior should catch the LinAlgError when peforming a
        # Cholesky decomposition and return NaN instead
        def potential_fn_chol(z):
            return torch.cholesky(z["K"])

        _, val = potential_grad(potential_fn_chol, z)
        self.assertTrue(torch.isnan(val))

        # Default behavior should not catch other errors
        def potential_fn_rterr_foo(z):
            raise RuntimeError("foo")

        with self.assertRaisesRegex(RuntimeError, "foo"):
            potential_grad(potential_fn_rterr_foo, z)

        # But once we register this specific error then it should
        def catch_runtime_error(e):
            return type(e) is RuntimeError and "foo" in str(e)

        register_exception_handler("foo_runtime", catch_runtime_error)
        _, val = potential_grad(potential_fn_rterr_foo, z)
        self.assertTrue(torch.isnan(val))

        # Unless the error message is different
        def potential_fn_rterr_bar(z):
            raise RuntimeError("bar")

        with self.assertRaisesRegex(RuntimeError, "bar"):
            potential_grad(potential_fn_rterr_bar, z)