jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def test_logsumexp_benchmark(size, axis, benchmark):
    X = at.matrix("X")
    X_max = at.max(X, axis=axis, keepdims=True)
    X_max = at.switch(at.isinf(X_max), 0, X_max)
    X_lse = at.log(at.sum(at.exp(X - X_max), axis=axis, keepdims=True)) + X_max

    rng = np.random.default_rng(23920)
    X_val = rng.normal(size=size)

    X_lse_fn = aesara.function([X], X_lse, mode="NUMBA")

    # JIT compile first
    _ = X_lse_fn(X_val)

    res = benchmark(X_lse_fn, X_val)

    exp_res = scipy.misc.logsumexp(X_val, axis=axis, keepdims=True)
    np.testing.assert_array_almost_equal(res, exp_res)
//  [python]
def test_logsumexp_benchmark(size, axis, benchmark):
    X = at.matrix("X")
    X_max = at.max(X, axis=axis, keepdims=True)
    X_max = at.switch(at.isinf(X_max), 0, X_max)
    X_lse = at.log(at.sum(at.exp(X - X_max), axis=axis, keepdims=True)) + X_max

    rng = np.random.default_rng(23920)
    X_val = rng.normal(size=size)

    X_lse_fn = aesara.function([X], X_lse, mode="NUMBA")

    # JIT compile first
    _ = X_lse_fn(X_val)

    res = benchmark(X_lse_fn, X_val)

    exp_res = scipy.special.logsumexp(X_val, axis=axis, keepdims=True)
    np.testing.assert_array_almost_equal(res, exp_res)
//  [python]
def fit_t0mu_gibbs(loc, t00_list, step, number_sample_zero, Tau, Sigma, mu, t00, b_mu, b_t0, TRIALS):
    def agg_NPE(t0):
        log_f = log_convolve_exp_norm(loc - t0, Tau, Sigma) - log_convolve_exp_norm(loc - t00_list, Tau, Sigma)
        return jit_agg_NPE(step, log_f, TRIALS)

    def t_t0(t0):
        nonlocal mu
        NPE, f_agg = agg_NPE(t0)
        NPE = np.insert(NPE, 0, 0)
        f_agg = np.insert(f_agg, 0, np.log(number_sample_zero))
        ans = opti.fmin_l_bfgs_b(lambda μ: μ - s.logsumexp(NPE * np.log(μ / mu) + f_agg), 
                                 x0=[mu], 
                                 approx_grad=True, 
                                 bounds=[b_mu], 
                                 maxfun=10000)
        mu = ans[0].item()
        return ans[1]

    t0_list = np.linspace(b_t0[0], b_t0[-1], 501)
    fval = np.array([t_t0(t_i) for t_i in t0_list])
    if np.std(fval) < 1e-4:
        t0 = np.mean(b_t0)
    else:
        t0 = t0_list[fval.argmin()]
    fval_new = t_t0(t0)
    return t0, mu
//  [python]
def fit_t0mu_gibbs(loc, t00_list, step, number_sample_zero, Tau, Sigma, mu, t00, b_mu, b_t0, TRIALS):
    def agg_NPE(t0):
        log_f = log_convolve_exp_norm(loc - t0, Tau, Sigma) - log_convolve_exp_norm(loc - t00_list, Tau, Sigma)
        return jit_agg_NPE(step, log_f, TRIALS)

    def t_t0(t0):
        nonlocal mu
        NPE, f_agg = agg_NPE(t0)
        NPE = np.insert(NPE, 0, 0)
        f_agg = np.insert(f_agg, 0, np.log(number_sample_zero))
        ans = opti.fmin_l_bfgs_b(lambda μ: μ - special.logsumexp(NPE * np.log(μ / mu) + f_agg), 
                                 x0=[mu], 
                                 approx_grad=True, 
                                 bounds=[b_mu], 
                                 maxfun=10000)
        mu = ans[0].item()
        return ans[1]

    t0_list = np.linspace(b_t0[0], b_t0[-1], 501)
    fval = np.array([t_t0(t_i) for t_i in t0_list])
    if np.std(fval) < 1e-4:
        t0 = np.mean(b_t0)
    else:
        t0 = t0_list[fval.argmin()]
    fval_new = t_t0(t0)
    return t0, mu
//  [python]
def __call__(self, theta):
        """
        Calculate loglikelihood(theta), as well as any derived parameters.

        N.B. this loglikelihood is not normalised.

        Parameters
        ----------
        theta: float or 1d numpy array

        Returns
        -------
        logl: float
            Loglikelihood
        phi: list of length nderived
            Any derived parameters.
        """
        assert theta.shape[0] % 2 == 0, (
            'ndim={} must be even'.format(theta.shape))
        # first component is
        # l(t) = 0.5 * loggamma(t - 10, 1, 1) + 0.5 * loggamma(t + 10, 1, 1)
        logl = np.log(0.5) + np.logaddexp(
            log_loggamma_pdf(theta[0] - 10, alpha=1, beta=1),
            log_loggamma_pdf(theta[0] + 10, alpha=1, beta=1))
        # second component is
        # l(t) = 0.5 * N(t - 10, 1, 1) + 0.5 * N(t + 10, 1, 1)
        logl += np.log(0.5) + np.logaddexp(
            log_gaussian_pdf(theta[1] - 10, sigma=1),
            log_gaussian_pdf(theta[1] + 10, sigma=1))
        if theta.shape[0] > 2:
            # remaining components are half loggamma and half gaussian
            bound_ind = (theta.shape[0] // 2) + 1
            logl += log_loggamma_pdf(theta[2:bound_ind], alpha=1, beta=1)
            logl += log_gaussian_pdf(theta[bound_ind:], sigma=1)
        return logl, [0.0] * self.nderived
//  [python]
def __call__(self, theta):
        """
        Calculate loglikelihood(theta), as well as any derived parameters.

        N.B. this loglikelihood is not normalised.

        Parameters
        ----------
        theta: float or 1d numpy array

        Returns
        -------
        logl: float
            Loglikelihood
        phi: list of length nderived
            Any derived parameters.
        """
        assert theta.shape[0] % 2 == 0, (
            'ndim={} must be even'.format(theta.shape))
        # first component is
        # l(t) = 0.5 * loggamma(t - 10, 1, 1) + 0.5 * loggamma(t + 10, 1, 1)
        logl = np.log(0.5) + scipy.special.logsumexp(
            [log_loggamma_pdf(th, alpha=1, beta=1) for th in
             [theta[0] - 10, theta[0] + 10]])
        # second component is
        # l(t) = 0.5 * N(t - 10, 1, 1) + 0.5 * N(t + 10, 1, 1)
        logl += np.log(0.5) + scipy.special.logsumexp(
            [log_gaussian_pdf(th, sigma=1) for th in
             [theta[1] - 10, theta[1] + 10]])
        if theta.shape[0] > 2:
            # remaining components are half loggamma and half gaussian
            bound_ind = (theta.shape[0] // 2) + 1
            logl += log_loggamma_pdf(theta[2:bound_ind], alpha=1, beta=1)
            logl += log_gaussian_pdf(theta[bound_ind:], sigma=1)
        return logl, [0.0] * self.nderived
//  [python]
def update(self, gradient_losses):
        """
        gradient_losses should be n * len(dt)
        """
        # add t
        self.TD += gradient_losses.shape[1]

        # add D
        self.TD += ((gradient_losses.shape[1] + 1) * gradient_losses.shape[1]) // 2

        # update estimated Z
        while self.Z < self.TD:
            self.Z *= 2

        lr = np.sqrt(self.lnN / self.Z)
        changes = lr * gradient_losses.sum(axis=1)
        log_multiple = np.log(self.w + 1e-14) - changes
        self.w = np.exp(log_multiple - s.logsumexp(log_multiple))
//  [python]
def update(self, gradient_losses):
        """
        gradient_losses should be n * len(dt)
        """
        # add t
        self.TD += gradient_losses.shape[1]

        # add D
        self.TD += ((gradient_losses.shape[1] + 1) * gradient_losses.shape[1]) // 2

        # update estimated Z
        while self.Z < self.TD:
            self.Z *= 2

        lr = np.sqrt(self.lnN / self.Z)
        changes = lr * gradient_losses.sum(axis=1)
        log_multiple = np.log(self.w + 1e-14) - changes
        self.w = np.exp(log_multiple - sc.logsumexp(log_multiple))
//  [python]
def fit_t0mu_guess(loc, step, number_sample_zero, Tau, Sigma, guess, mu, t00, b_mu, b_t0, TRIALS):
    def agg_NPE(t0):
        log_f = log_convolve_exp_norm(loc - t0, Tau, Sigma) + guess
        return jit_agg_NPE(step, log_f, TRIALS)

    def t_t0(t0):
        nonlocal mu
        NPE, f_agg = agg_NPE(t0)
        if number_sample_zero != 0:
            NPE = np.insert(NPE, 0, 0)
            f_agg = np.insert(f_agg, 0, np.log(number_sample_zero))
        ans = opti.fmin_l_bfgs_b(lambda μ: μ - s.logsumexp(NPE * np.log(μ / mu) + f_agg), 
                                 x0=[mu], 
                                 approx_grad=True, 
                                 bounds=[b_mu], 
                                 maxfun=10000)
        mu = ans[0].item()
        return ans[1]

    # ans = opti.fmin_l_bfgs_b(t_t0, x0=[t00], approx_grad=True, bounds=[b_t0], maxfun=10000)
    # t0 = ans[0].item()
    # fval_old = ans[1].item()

    t0_list = np.linspace(b_t0[0], b_t0[-1], 501)
    fval = np.array([t_t0(t_i) for t_i in t0_list])
    if np.std(fval) < 1e-4:
        t0 = np.mean(b_t0)
    else:
        t0 = t0_list[fval.argmin()]
    fval_new = t_t0(t0)
    return t0, mu
//  [python]
def fit_t0mu_guess(loc, step, number_sample_zero, Tau, Sigma, guess, mu, t00, b_mu, b_t0, TRIALS):
    def agg_NPE(t0):
        log_f = log_convolve_exp_norm(loc - t0, Tau, Sigma) + guess
        return jit_agg_NPE(step, log_f, TRIALS)

    def t_t0(t0):
        nonlocal mu
        NPE, f_agg = agg_NPE(t0)
        if number_sample_zero != 0:
            NPE = np.insert(NPE, 0, 0)
            f_agg = np.insert(f_agg, 0, np.log(number_sample_zero))
        ans = opti.fmin_l_bfgs_b(lambda μ: μ - special.logsumexp(NPE * np.log(μ / mu) + f_agg), 
                                 x0=[mu], 
                                 approx_grad=True, 
                                 bounds=[b_mu], 
                                 maxfun=10000)
        mu = ans[0].item()
        return ans[1]

    # ans = opti.fmin_l_bfgs_b(t_t0, x0=[t00], approx_grad=True, bounds=[b_t0], maxfun=10000)
    # t0 = ans[0].item()
    # fval_old = ans[1].item()

    t0_list = np.linspace(b_t0[0], b_t0[-1], 501)
    fval = np.array([t_t0(t_i) for t_i in t0_list])
    if np.std(fval) < 1e-4:
        t0 = np.mean(b_t0)
    else:
        t0 = t0_list[fval.argmin()]
    fval_new = t_t0(t0)
    return t0, mu
//  [python]
def generalized_pg_mixture_2nd(k, alphas, betas):
    
    iters=[numpy.array(i) for i in bars_and_stars_iterator(int(k), len(betas))]
    
    log_res=[]
    for it in iters:
        
        log_res.append(calc_pg(it, alphas, betas).sum())
    
    return s.logsumexp(log_res)
//  [python]
def generalized_pg_mixture_2nd(k, alphas, betas):
    
    iters=[numpy.array(i) for i in bars_and_stars_iterator(int(k), len(betas))]
    
    log_res=[]
    for it in iters:
        
        log_res.append(calc_pg(it, alphas, betas).sum())
    
    return scipy.special.logsumexp(log_res)
//  [python]
def compute_log_likelihood(self, log2_ptr, binned_counts):
        """Compute the log likelihood of reordered bin probabilities given a log2(PTR).

        Parameters
        ----------
            log2_ptr : float
                The Log2(PTR)
            binned_counts : np.array
                Binned read counts along the genome

        Returns
        -------
            log_lk : float
                The log likelihood
        """
        # origin and terminus are at positions 0 and 1
        ori_loc = 0
        ter_loc = 1

        x = np.linspace(0, 1, binned_counts.size)
        # change to base e
        loge_ptr = log2_ptr / np.log2(np.e)
        log_probs = np.array([loge_ptr * xi for xi in x])
        log_norm = np.log(np.sum(np.exp(log_probs)))
        log_probs -= log_norm

        # x = np.linspace(0, 1, binned_counts.size)
        # # compute normalized probabilities
        # log_probs = np.array([log2_ptr*xi for xi in x])
        # log_norm = np.log2(np.power(2, log_probs).sum())
        # log_probs -= log_norm

        log_lk = (binned_counts * log_probs).sum()

        return log_lk
//  [python]
def compute_log_likelihood(self, log2_ptr, binned_counts):
        """Compute the log likelihood of reordered bin probabilities given a log2(PTR).

        Parameters
        ----------
            log2_ptr : float
                The Log2(PTR)
            binned_counts : np.array
                Binned read counts along the genome

        Returns
        -------
            log_lk : float
                The log likelihood
        """
        # origin and terminus are at positions 0 and 1
        ori_loc = 0
        ter_loc = 1

        x = np.linspace(0, 1, binned_counts.size)
        # change to base e
        loge_ptr = log2_ptr / np.log2(np.e)
        log_probs = np.array([loge_ptr * xi for xi in x])
        log_norm = scipy.special.logsumexp(log_probs)
        log_probs -= log_norm

        # x = np.linspace(0, 1, binned_counts.size)
        # # compute normalized probabilities
        # log_probs = np.array([log2_ptr*xi for xi in x])
        # log_norm = np.log2(np.power(2, log_probs).sum())
        # log_probs -= log_norm

        log_lk = (binned_counts * log_probs).sum()

        return log_lk
//  [python]
def process_run(self):
        """Use the running totals of realspace running_sum and nsamples_per_chain to
        calculate an estimate of the inverse evidence, its variance, and the
        variance of the variance.

        This method is ran each time chains are added to update the inverse
        variance estimates from the running totals.

        """
        nsamples_per_chain = self.nsamples_per_chain

        evidence_inv = jnp.sum(self.running_sum)
        nsamples = jnp.sum(self.nsamples_per_chain)

        evidence_inv /= nsamples

        self.ln_evidence_inv_per_chain = (
            jnp.log(self.running_sum) - self.shift_value - jnp.log(nsamples_per_chain)
        )

        """
        The following code computes the exponents of the variance and variance 
        of the variance where possible in log space to avoid overflow errors.

        This is a log-space representation of the real-space statistics. One may
        alternatively compute the log-space statistics, but should note that 
        simply taking the exponential of log-space statistics is NOT the same as 
        computing the real-space statistics.
        """

        # Precompute differential vectors.
        z_i = np.abs((self.running_sum / self.nsamples_per_chain) - evidence_inv)
        y_i = z_i * (self.nsamples_per_chain ** (0.25))
        z_i *= self.nsamples_per_chain ** (0.5)

        # Compute exponents using logsumexp for numerical stability.
        evidence_inv_var_ln_temp = s.logsumexp(2.0 * np.log(z_i)) - np.log(nsamples)
        kur_ln = (
            s.logsumexp(4.0 * np.log(y_i))
            - np.log(nsamples)
            - 2.0 * evidence_inv_var_ln_temp
        )
        kur = np.exp(kur_ln)
        self.kurtosis = kur
        self.ln_kurtosis = kur_ln

        # Compute effective chain lengths.
        n_eff = jnp.sum(jnp.square(nsamples_per_chain))
        n_eff = nsamples * nsamples / n_eff
        self.n_eff = n_eff

        # Compute inverse evidence values as a log-space representation of the
        # real-space statistics to attempt to avoid float overflow.
        self.ln_evidence_inv = np.log(evidence_inv) - self.shift_value
        self.ln_evidence_inv_var = (
            evidence_inv_var_ln_temp - 2 * self.shift_value - np.log(n_eff - 1)
        )
        self.ln_evidence_inv_var_var = (
            2.0 * evidence_inv_var_ln_temp
            - 3.0 * np.log(n_eff)
            - 4.0 * self.shift_value
            + np.log((kur - 1) + 2.0 / (n_eff - 1))
        )

        return
//  [python]
def process_run(self):
        """Use the running totals of realspace running_sum and nsamples_per_chain to
        calculate an estimate of the inverse evidence, its variance, and the
        variance of the variance.

        This method is ran each time chains are added to update the inverse
        variance estimates from the running totals.

        """
        nsamples_per_chain = self.nsamples_per_chain

        evidence_inv = jnp.sum(self.running_sum)
        nsamples = jnp.sum(self.nsamples_per_chain)

        evidence_inv /= nsamples

        self.ln_evidence_inv_per_chain = (
            jnp.log(self.running_sum) - self.shift_value - jnp.log(nsamples_per_chain)
        )

        """
        The following code computes the exponents of the variance and variance 
        of the variance where possible in log space to avoid overflow errors.

        This is a log-space representation of the real-space statistics. One may
        alternatively compute the log-space statistics, but should note that 
        simply taking the exponential of log-space statistics is NOT the same as 
        computing the real-space statistics.
        """

        # Precompute differential vectors.
        z_i = np.abs((self.running_sum / self.nsamples_per_chain) - evidence_inv)
        y_i = z_i * (self.nsamples_per_chain ** (0.25))
        z_i *= self.nsamples_per_chain ** (0.5)

        # Compute exponents using logsumexp for numerical stability.
        evidence_inv_var_ln_temp = sp.logsumexp(2.0 * np.log(z_i)) - np.log(nsamples)
        kur_ln = (
            sp.logsumexp(4.0 * np.log(y_i))
            - np.log(nsamples)
            - 2.0 * evidence_inv_var_ln_temp
        )
        kur = np.exp(kur_ln)
        self.kurtosis = kur
        self.ln_kurtosis = kur_ln

        # Compute effective chain lengths.
        n_eff = jnp.sum(jnp.square(nsamples_per_chain))
        n_eff = nsamples * nsamples / n_eff
        self.n_eff = n_eff

        # Compute inverse evidence values as a log-space representation of the
        # real-space statistics to attempt to avoid float overflow.
        self.ln_evidence_inv = np.log(evidence_inv) - self.shift_value
        self.ln_evidence_inv_var = (
            evidence_inv_var_ln_temp - 2 * self.shift_value - np.log(n_eff - 1)
        )
        self.ln_evidence_inv_var_var = (
            2.0 * evidence_inv_var_ln_temp
            - 3.0 * np.log(n_eff)
            - 4.0 * self.shift_value
            + np.log((kur - 1) + 2.0 / (n_eff - 1))
        )

        return
//  [python]
def estimate_ptr_maximum_likelihood(self, binned_counts):
        """Compute maximum likelihood estimates of log2(PTR) gived observed reads.

        Parameters
        ----------
            binned_counts : np.array
                Binned read counts for a genome

        Returns
        -------
            log2_ptr : float
                The maximum likelihood log2(PTR) estimate
            intercept : float
                Probability a reads is in the leftmost bin
            log_lk : float
                The log likelihood

        """
        f = lambda x: -self.compute_log_likelihood(x, binned_counts)

        log2_ptr, fun, it, funcalls = scipy.optimize.brent(f, full_output=True)
        log_lk = -fun

        x = np.linspace(0, 1, binned_counts.size)
        # compute normalized probabilities
        loge_ptr = log2_ptr / np.log2(np.e)
        log_probs = np.array([loge_ptr * xi for xi in x])
        log_norm = np.max(log_probs) + np.log(np.sum(np.exp(log_probs - np.max(log_probs))))
        log_probs -= log_norm
        probs = np.exp(log_probs)
        intercept = probs[0]

        # the orientation of the bins is arbitrary,
        # so let's orient with a positive slope
        flip = False
        if log2_ptr < 0:
            intercept = log2_ptr + intercept
            log2_ptr = np.abs(log2_ptr)
            flip = True

        return log2_ptr, intercept, log_lk, flip
//  [python]
def estimate_ptr_maximum_likelihood(self, binned_counts):
        """Compute maximum likelihood estimates of log2(PTR) gived observed reads.

        Parameters
        ----------
            binned_counts : np.array
                Binned read counts for a genome

        Returns
        -------
            log2_ptr : float
                The maximum likelihood log2(PTR) estimate
            intercept : float
                Probability a reads is in the leftmost bin
            log_lk : float
                The log likelihood

        """
        f = lambda x: -self.compute_log_likelihood(x, binned_counts)

        log2_ptr, fun, it, funcalls = scipy.optimize.brent(f, full_output=True)
        log_lk = -fun

        x = np.linspace(0, 1, binned_counts.size)
        # compute normalized probabilities
        loge_ptr = log2_ptr / np.log2(np.e)
        log_probs = np.array([loge_ptr * xi for xi in x])
        log_norm = scipy.special.logsumexp(log_probs)
        log_probs -= log_norm
        probs = np.exp(log_probs)
        intercept = probs[0]

        # the orientation of the bins is arbitrary,
        # so let's orient with a positive slope
        flip = False
        if log2_ptr < 0:
            intercept = log2_ptr + intercept
            log2_ptr = np.abs(log2_ptr)
            flip = True

        return log2_ptr, intercept, log_lk, flip
//  [python]
def pg_log_python(k, weights, alpha_individual=0.0, extra_prior_counter=0.0, extra_prior_weight=1.0):
    
    log_deltas=[0.0]
    log_inner_factors=[]
    log_weight_prefactors=0.0

    log_weight_prefactors+=-((1.0+alpha_individual)*numpy.log(1.0+weights)).sum()

    first_fac=1.0+alpha_individual
    
    log_first_fac= numpy.log(first_fac)
    log_first_var=-numpy.log(1.0+1.0/weights)

    running_factor_vec=0.0

    if(k>0):
        
        for i in (numpy.arange(k)+1):
            
            running_factor_vec+=log_first_var
            
            # summing assuming all summands positive .. which is the only well-defined region here
            
            res=numpy.logaddexp.reduce(log_first_fac+running_factor_vec)
            
            log_inner_factors.append(res)
            new_delta=numpy.logaddexp.reduce( numpy.array(log_inner_factors[::-1])+numpy.array(log_deltas))-numpy.log(i)
            log_deltas.append(new_delta)
 
            #log_inner_factors.append((first_fac*(first_var**i)-second_fac*(second_var**i)).sum())
            #log_deltas.append( (numpy.array(inner_factors)*numpy.array(deltas[::-1])).sum()/float(i))
            
    return log_deltas[-1]+log_weight_prefactors
//  [python]
def pg_log_python(k, weights, alpha_individual=0.0, extra_prior_counter=0.0, extra_prior_weight=1.0):
    
    log_deltas=[0.0]
    log_inner_factors=[]
    log_weight_prefactors=0.0

    log_weight_prefactors+=-((1.0+alpha_individual)*numpy.log(1.0+weights)).sum()

    first_fac=1.0+alpha_individual
    
    log_first_fac= numpy.log(first_fac)
    log_first_var=-numpy.log(1.0+1.0/weights)

    running_factor_vec=0.0

    if(k>0):
        
        for i in (numpy.arange(k)+1):
            
            running_factor_vec+=log_first_var
            
            # summing assuming all summands positive .. which is the only well-defined region here
            
            res=scipy.special.logsumexp(log_first_fac+running_factor_vec)
            
            log_inner_factors.append(res)
            new_delta=scipy.special.logsumexp( numpy.array(log_inner_factors[::-1])+numpy.array(log_deltas))-numpy.log(i)
            log_deltas.append(new_delta)
 
            #log_inner_factors.append((first_fac*(first_var**i)-second_fac*(second_var**i)).sum())
            #log_deltas.append( (numpy.array(inner_factors)*numpy.array(deltas[::-1])).sum()/float(i))
            
    return log_deltas[-1]+log_weight_prefactors
//  [python]
def get_logw(ns_run, simulate=False):
    r"""Calculates the log posterior weights of the samples (using logarithms
    to avoid overflow errors with very large or small values).

    Uses the trapezium rule such that the weight of point i is

    .. math:: w_i = \mathcal{L}_i (X_{i-1} - X_{i+1}) / 2

    Parameters
    ----------
    ns_run: dict
        Nested sampling run dict (see data_processing module docstring for more
        details).
    simulate: bool, optional
        Should log prior volumes logx be simulated from their distribution (if
        false their expected values are used).

    Returns
    -------
    logw: 1d numpy array
        Log posterior masses of points.
    """
    try:
        # find logX value for each point
        logx = get_logx(ns_run['nlive_array'], simulate=simulate)
        logw = np.zeros(ns_run['logl'].shape[0])
        # Vectorized trapezium rule: w_i prop to (X_{i-1} - X_{i+1}) / 2
        logw[1:-1] = log_subtract(logx[:-2], logx[2:]) - np.log(2)
        # Assign all prior volume closest to first point X_first to that point:
        # that is from logx=0 to logx=log((X_first + X_second) / 2)
        logw[0] = log_subtract(0, np.logaddexp(logx[0], logx[1]) - np.log(2))
        # Assign all prior volume closest to final point X_last to that point:
        # that is from logx=log((X_penultimate + X_last) / 2) to logx=-inf
        logw[-1] = np.logaddexp(logx[-2], logx[-1]) - np.log(2)
        # multiply by likelihood (add in log space)
        logw += ns_run['logl']
        return logw
    except IndexError:
        if ns_run['logl'].shape[0] == 1:
            # If there is only one point in the run then assign all prior
            # volume X \in (0, 1) to that point, so the weight is just
            # 1 * logl_0 = logl_0
            return copy.deepcopy(ns_run['logl'])
        else:
            raise
//  [python]
def get_logw(ns_run, simulate=False):
    r"""Calculates the log posterior weights of the samples (using logarithms
    to avoid overflow errors with very large or small values).

    Uses the trapezium rule such that the weight of point i is

    .. math:: w_i = \mathcal{L}_i (X_{i-1} - X_{i+1}) / 2

    Parameters
    ----------
    ns_run: dict
        Nested sampling run dict (see data_processing module docstring for more
        details).
    simulate: bool, optional
        Should log prior volumes logx be simulated from their distribution (if
        false their expected values are used).

    Returns
    -------
    logw: 1d numpy array
        Log posterior masses of points.
    """
    try:
        # find logX value for each point
        logx = get_logx(ns_run['nlive_array'], simulate=simulate)
        logw = np.zeros(ns_run['logl'].shape[0])
        # Vectorized trapezium rule: w_i prop to (X_{i-1} - X_{i+1}) / 2
        logw[1:-1] = log_subtract(logx[:-2], logx[2:]) - np.log(2)
        # Assign all prior volume closest to first point X_first to that point:
        # that is from logx=0 to logx=log((X_first + X_second) / 2)
        logw[0] = log_subtract(0, scipy.special.logsumexp([logx[0], logx[1]]) -
                               np.log(2))
        # Assign all prior volume closest to final point X_last to that point:
        # that is from logx=log((X_penultimate + X_last) / 2) to logx=-inf
        logw[-1] = scipy.special.logsumexp([logx[-2], logx[-1]]) - np.log(2)
        # multiply by likelihood (add in log space)
        logw += ns_run['logl']
        return logw
    except IndexError:
        if ns_run['logl'].shape[0] == 1:
            # If there is only one point in the run then assign all prior
            # volume X \in (0, 1) to that point, so the weight is just
            # 1 * logl_0 = logl_0
            return copy.deepcopy(ns_run['logl'])
        else:
            raise
//  [python]
def generate_log_stirling(max_val=1000):

    arr=numpy.zeros(shape=(max_val+1, max_val+1))*(-numpy.inf)
    arr[0][0]=0.0
    for i in range(max_val+1):
        for j in range(i+1):
            if(i==0 and j==0):
                arr[i][j]=0.0
                continue
            
            if(j==0):
                arr[i][j]=-numpy.inf
                continue
            if(i==j):
                arr[i][j]=0.0
                continue
          
            arr[i][j]=scipy.special.logsumexp([numpy.log(i-1.0)+arr[i-1][j],arr[i-1][j-1]])
    
    return arr
//  [python]
def log_normalization_constant(self, all_g):
        """
        Computes normalization constant for log p(G), i.e. `Z = log(sum_G p(g))`
        all_g: list of igraph.Graph objects
        Returns:
            float
        """
        log_prob_g_unn = np.zeros(len(all_g))
        for i, g in enumerate(
            tqdm.tqdm(
                all_g, desc="p(G) log_normalization_constant", disable=not self.verbose
            )
        ):
            log_prob_g_unn[i] = self.unnormalized_log_prob(g=g)
        log_prob_sum_g = scipy.misc.logsumexp(log_prob_g_unn)
        return log_prob_sum_g
//  [python]
def log_normalization_constant(self, all_g):
        """
        Computes normalization constant for log p(G), i.e. `Z = log(sum_G p(g))`
        all_g: list of igraph.Graph objects
        Returns:
            float
        """
        log_prob_g_unn = np.zeros(len(all_g))
        for i, g in enumerate(
            tqdm.tqdm(
                all_g, desc="p(G) log_normalization_constant", disable=not self.verbose
            )
        ):
            log_prob_g_unn[i] = self.unnormalized_log_prob(g=g)
        log_prob_sum_g = scipy.special.logsumexp(log_prob_g_unn)
        return log_prob_sum_g
//  [python]
def __init__(self, n_vars, gibbs_temp=10.0, sparsity_factor=0.0, verbose=False):
        super(GibbsDAGDistributionFull, self).__init__(n_vars=n_vars, verbose=verbose)
        assert (
            n_vars <= 4
        ), "Cannot use this for higher dimensional variables, Try UniformDAGDistributionRejection instead"
        self.n_vars = n_vars
        self.verbose = verbose
        self.gibbs_temp = gibbs_temp
        self.sparsity_factor = sparsity_factor
        all_g = all_combinations(
            n_vars, return_adj=True
        )  # Do not store this in interest of memory
        dagness = np.zeros(len(all_g))
        for i, j in enumerate(all_g):
            dagness[i] = expm_np(j, self.n_vars)
        self.logits = -gibbs_temp * dagness - sparsity_factor * np.sum(
            all_g, axis=(-1, -2)
        )
        self.z_g = s.logsumexp(self.logits)
//  [python]
def __init__(self, n_vars, gibbs_temp=10.0, sparsity_factor=0.0, verbose=False):
        super(GibbsDAGDistributionFull, self).__init__(n_vars=n_vars, verbose=verbose)
        assert (
            n_vars <= 4
        ), "Cannot use this for higher dimensional variables, Try UniformDAGDistributionRejection instead"
        self.n_vars = n_vars
        self.verbose = verbose
        self.gibbs_temp = gibbs_temp
        self.sparsity_factor = sparsity_factor
        all_g = all_combinations(
            n_vars, return_adj=True
        )  # Do not store this in interest of memory
        dagness = np.zeros(len(all_g))
        for i, j in enumerate(all_g):
            dagness[i] = expm_np(j, self.n_vars)
        self.logits = -gibbs_temp * dagness - sparsity_factor * np.sum(
            all_g, axis=(-1, -2)
        )
        self.z_g = scipy.special.logsumexp(self.logits)
//  [python]
def pgpg_log_python(k, weights, mean_adjustment):
    
    ## fix extra adjustment in standard PG to 0
    alpha_individual=0.0
    
    log_weight_prefactors=-(alpha_individual*numpy.log(1.0+weights)).sum()
   
    Cs=1.0/(2+2*weights)

    log_weight_prefactors+=-((1.0+mean_adjustment)*numpy.log(2.0-2*Cs)).sum()
    
    log_E_s=-numpy.log(1.+1./weights)
    #one_minus_c=(1+2*weights)/(2+2*weights)
    
    first_fac=1.0+mean_adjustment
    signs_first=numpy.where(first_fac>0, 1.0, -1.0)
    log_first_fac=numpy.where(signs_first>0, numpy.log(first_fac), numpy.log(-first_fac))
    
    second_fac=-(1.0+mean_adjustment-alpha_individual)
    signs_second=numpy.where(second_fac>0, 1.0, -1.0)
    log_second_fac=numpy.where(signs_second>0, numpy.log(second_fac), numpy.log(-second_fac))

    log_first_var=log_E_s-numpy.log(1.0-Cs)
    log_second_var=log_E_s

    log_deltas=[0.0]
    log_inner_factors=[]
    
    running_factor_vec_first=0.0
    running_factor_vec_second=0.0
    
    if(k>0):
        
        for i in (numpy.arange(k)+1):
            
            running_factor_vec_first+=log_first_var
            running_factor_vec_second+=log_second_var
            
            sum1=numpy.max(log_first_fac+running_factor_vec_first) + numpy.log(numpy.sum(signs_first*numpy.exp(log_first_fac+running_factor_vec_first-numpy.max(log_first_fac+running_factor_vec_first))))
            sum2=numpy.max(log_second_fac+running_factor_vec_second) + numpy.log(numpy.sum(signs_second*numpy.exp(log_second_fac+running_factor_vec_second-numpy.max(log_second_fac+running_factor_vec_second))))

            sign1=numpy.sign(numpy.sum(signs_first*numpy.exp(log_first_fac+running_factor_vec_first)))
            sign2=numpy.sign(numpy.sum(signs_second*numpy.exp(log_second_fac+running_factor_vec_second)))

            res=numpy.max([sum1, sum2]) + numpy.log(numpy.sum([sign1*numpy.exp(sum1), sign2*numpy.exp(sum2)]))


            log_inner_factors.append(res)
            new_delta=numpy.max(numpy.array(log_inner_factors[::-1])+numpy.array(log_deltas)) + numpy.log(numpy.sum(numpy.exp(numpy.array(log_inner_factors[::-1])+numpy.array(log_deltas)-numpy.max(numpy.array(log_inner_factors[::-1])+numpy.array(log_deltas)))))-numpy.log(i)
            log_deltas.append(new_delta)
 
            #log_inner_factors.append((first_fac*(first_var**i)-second_fac*(second_var**i)).sum())
            #log_deltas.append( (numpy.array(inner_factors)*numpy.array(deltas[::-1])).sum()/float(i))
            
    
    return log_deltas[-1]+log_weight_prefactors
//  [python]
def pgpg_log_python(k, weights, mean_adjustment):
    
    ## fix extra adjustment in standard PG to 0
    alpha_individual=0.0
    
    log_weight_prefactors=-(alpha_individual*numpy.log(1.0+weights)).sum()
   
    Cs=1.0/(2+2*weights)

    log_weight_prefactors+=-((1.0+mean_adjustment)*numpy.log(2.0-2*Cs)).sum()
    
    log_E_s=-numpy.log(1.+1./weights)
    #one_minus_c=(1+2*weights)/(2+2*weights)
    
    first_fac=1.0+mean_adjustment
    signs_first=numpy.where(first_fac>0, 1.0, -1.0)
    log_first_fac=numpy.where(signs_first>0, numpy.log(first_fac), numpy.log(-first_fac))
    
    second_fac=-(1.0+mean_adjustment-alpha_individual)
    signs_second=numpy.where(second_fac>0, 1.0, -1.0)
    log_second_fac=numpy.where(signs_second>0, numpy.log(second_fac), numpy.log(-second_fac))

    log_first_var=log_E_s-numpy.log(1.0-Cs)
    log_second_var=log_E_s

    log_deltas=[0.0]
    log_inner_factors=[]
    
    running_factor_vec_first=0.0
    running_factor_vec_second=0.0
    
    if(k>0):
        
        for i in (numpy.arange(k)+1):
            
            running_factor_vec_first+=log_first_var
            running_factor_vec_second+=log_second_var
            
            sum1,sign1=scipy.special.logsumexp(log_first_fac+running_factor_vec_first, b=signs_first,return_sign=True)
            sum2,sign2=scipy.special.logsumexp(log_second_fac+running_factor_vec_second, b=signs_second,return_sign=True)
            
            res=scipy.special.logsumexp([sum1, sum2], b=[sign1,sign2])
            
            log_inner_factors.append(res)
            new_delta=scipy.special.logsumexp( numpy.array(log_inner_factors[::-1])+numpy.array(log_deltas))-numpy.log(i)
            log_deltas.append(new_delta)
 
            #log_inner_factors.append((first_fac*(first_var**i)-second_fac*(second_var**i)).sum())
            #log_deltas.append( (numpy.array(inner_factors)*numpy.array(deltas[::-1])).sum()/float(i))
            
    
    return log_deltas[-1]+log_weight_prefactors
//  [python]
def sample_mixture_gaussians_2d(w, mu, L, n_samples):
    """
    Given k mixtures' parameters, sample
    :param w: mixing probabilities summing to 1
    :param mu: k, 2
    :param L: k, 2, 2 | LL^T = var
        - so (L^-1)' L^-1 = var^-1
        - hence ||L^-1(x-mu)||^2 = (x-mu)'var^-1(x-mu)
    :param n_samples:
    :return:
        inds: n_samples,
        x: n_samples, 2
        nll: n_samples | for importance sampling
    """
    L_inv = make_L_inv_2d(L)
    det_var = (L[:, 0, 0] * L[:, 1, 1]) ** 2

    inds = np.random.choice(w.size, n_samples, p=w)
    noise = np.random.randn(2, n_samples)
    x = mu[inds] + np.einsum('kij,jk->ki', L[inds], noise)

    dif = x[:, np.newaxis] - mu[np.newaxis]  # n_samples, k, 2
    z = np.einsum('kij,skj->ski', L_inv, dif)  # n_samples, k, 2
    nlog_p = .5 * ((z ** 2).sum(axis=2) + np.log(det_var))  # n_samples, k
    # nll = -np.log((np.exp(-log_p) * w).sum(axis=1))
    nll = -s.logsumexp(-nlog_p, b=w, axis=1)
    return inds, x, nll
//  [python]
def sample_mixture_gaussians_2d(w, mu, L, n_samples):
    """
    Given k mixtures' parameters, sample
    :param w: mixing probabilities summing to 1
    :param mu: k, 2
    :param L: k, 2, 2 | LL^T = var
        - so (L^-1)' L^-1 = var^-1
        - hence ||L^-1(x-mu)||^2 = (x-mu)'var^-1(x-mu)
    :param n_samples:
    :return:
        inds: n_samples,
        x: n_samples, 2
        nll: n_samples | for importance sampling
    """
    L_inv = make_L_inv_2d(L)
    det_var = (L[:, 0, 0] * L[:, 1, 1]) ** 2

    inds = np.random.choice(w.size, n_samples, p=w)
    noise = np.random.randn(2, n_samples)
    x = mu[inds] + np.einsum('kij,jk->ki', L[inds], noise)

    dif = x[:, np.newaxis] - mu[np.newaxis]  # n_samples, k, 2
    z = np.einsum('kij,skj->ski', L_inv, dif)  # n_samples, k, 2
    nlog_p = .5 * ((z ** 2).sum(axis=2) + np.log(det_var))  # n_samples, k
    # nll = -np.log((np.exp(-log_p) * w).sum(axis=1))
    nll = -spe.logsumexp(-nlog_p, b=w, axis=1)
    return inds, x, nll
//  [python]
def dist_from_logprobs(value_to_unnorm_logprob):
    """An attempt to reuse the current __init__ code above, which works in prob
    space, by converting up from logprob input.  Not sure yet this is the right
    thing to do -- maybe it would be better to use log-domain arithmetic to do
    all the code in log-domain.
    
    input to this function is a dictionary.
    """
    unnorm_logprobs = np.array(list(value_to_unnorm_logprob.values()))
    logZ = scipy.misc.logsumexp(unnorm_logprobs)
    value_to_norm_logprobs = {k: ulp-logZ for k,ulp in list(value_to_unnorm_logprob.items())}
    # Danger step!  could get unjustified 0s.
    value_to_norm_probs = {k: np.exp(lp) for k,lp in list(value_to_norm_logprobs.items())}
    # If logsumexp is implemented well, the prob sum should now be 1.0.
    # (I actually don't know if floating-point semantics allow it to be the
    # case that that can be guaranteed.)
    # We could pass in normalizer=1.0 to the constructor to save a bit of time.
    # But just in case we'll let it recompute it.
    return CategDist(value_to_norm_probs)
//  [python]
def dist_from_logprobs(value_to_unnorm_logprob):
    """An attempt to reuse the current __init__ code above, which works in prob
    space, by converting up from logprob input.  Not sure yet this is the right
    thing to do -- maybe it would be better to use log-domain arithmetic to do
    all the code in log-domain.
    
    input to this function is a dictionary.
    """
    unnorm_logprobs = np.array(list(value_to_unnorm_logprob.values()))
    #logZ = scipy.misc.logsumexp(unnorm_logprobs) #deprecated, use new logsumexp
    logZ = scipy.special.logsumexp(unnorm_logprobs)
    value_to_norm_logprobs = {k: ulp-logZ for k,ulp in list(value_to_unnorm_logprob.items())}
    # Danger step!  could get unjustified 0s.
    value_to_norm_probs = {k: np.exp(lp) for k,lp in list(value_to_norm_logprobs.items())}
    # If logsumexp is implemented well, the prob sum should now be 1.0.
    # (I actually don't know if floating-point semantics allow it to be the
    # case that that can be guaranteed.)
    # We could pass in normalizer=1.0 to the constructor to save a bit of time.
    # But just in case we'll let it recompute it.
    return CategDist(value_to_norm_probs)
//  [python]
def logsumexp(a, axis=None, b=None, keepdims=False, return_sign=False):
    r"""Compute the log of the sum of exponentials of input elements.

    This function has the same call signature as
    :func:`scipy.special.logsumexp` and mirrors scipy's behaviour except for
    ``-np.inf`` input. If a and b are both ``-inf`` then scipy's function will
    output ``nan`` whereas here we use:

    .. math::

        \lim_{x \to -\infty} x \exp(x) = 0

    Thus, if ``a=-inf`` in ``log(sum(b * exp(a))`` then we can set ``b=0``
    such that that term is ignored in the sum.
    """
    if b is None:
        b = np.ones_like(a)
    b = np.where(a == -np.inf, 0, b)
    return s.logsumexp(a, axis=axis, b=b, keepdims=keepdims,
                             return_sign=return_sign)
//  [python]
def logsumexp(a, axis=None, b=None, keepdims=False, return_sign=False):
    r"""Compute the log of the sum of exponentials of input elements.

    This function has the same call signature as
    :func:`scipy.special.logsumexp` and mirrors scipy's behaviour except for
    ``-np.inf`` input. If a and b are both ``-inf`` then scipy's function will
    output ``nan`` whereas here we use:

    .. math::

        \lim_{x \to -\infty} x \exp(x) = 0

    Thus, if ``a=-inf`` in ``log(sum(b * exp(a))`` then we can set ``b=0``
    such that that term is ignored in the sum.
    """
    if b is None:
        b = np.ones_like(a)
    b = np.where(a == -np.inf, 0, b)
    return special.logsumexp(a, axis=axis, b=b, keepdims=keepdims,
                             return_sign=return_sign)
//  [python]
def sample_mixture_gaussians_nd(w, mu, L, n_samples, L_inv=None):
    """
    Given k mixtures' parameters, sample
    :param w: mixing probabilities summing to 1
    :param mu: k, n | k n-dimensional gaussian mixture components
    :param L: k, n, n | LL^T = var
        - so (L^-1)' L^-1 = var^-1
        - hence ||L^-1(x-mu)||^2 = (x-mu)'var^-1(x-mu)
    :param n_samples:
    :return:
        inds: n_samples,
        x: n_samples, n
        nll: n_samples | for importance sampling
    """
    if L_inv is None:
        L_inv = make_L_inv_nd(L)
    n = L.shape[1]
    log_det_var = 2 * np.log(L[:, np.arange(n), np.arange(n)]).sum(axis=1)  # k,

    inds = np.random.choice(w.size, n_samples, p=w)
    noise = np.random.randn(n, n_samples)
    x = mu[inds] + np.einsum('kij,jk->ki', L[inds], noise)

    dif = x[:, np.newaxis] - mu[np.newaxis]  # n_samples, k, n
    z = np.einsum('kij,skj->ski', L_inv, dif)  # n_samples, k, n
    nlog_p = .5 * ((z ** 2).sum(axis=2) + log_det_var)  # n_samples, k
    # nll = -np.log((np.exp(-log_p) * w).sum(axis=1))
    nll = -s.logsumexp(-nlog_p, b=w, axis=1)
    return inds, x, nll
//  [python]
def sample_mixture_gaussians_nd(w, mu, L, n_samples, L_inv=None):
    """
    Given k mixtures' parameters, sample
    :param w: mixing probabilities summing to 1
    :param mu: k, n | k n-dimensional gaussian mixture components
    :param L: k, n, n | LL^T = var
        - so (L^-1)' L^-1 = var^-1
        - hence ||L^-1(x-mu)||^2 = (x-mu)'var^-1(x-mu)
    :param n_samples:
    :return:
        inds: n_samples,
        x: n_samples, n
        nll: n_samples | for importance sampling
    """
    if L_inv is None:
        L_inv = make_L_inv_nd(L)
    n = L.shape[1]
    log_det_var = 2 * np.log(L[:, np.arange(n), np.arange(n)]).sum(axis=1)  # k,

    inds = np.random.choice(w.size, n_samples, p=w)
    noise = np.random.randn(n, n_samples)
    x = mu[inds] + np.einsum('kij,jk->ki', L[inds], noise)

    dif = x[:, np.newaxis] - mu[np.newaxis]  # n_samples, k, n
    z = np.einsum('kij,skj->ski', L_inv, dif)  # n_samples, k, n
    nlog_p = .5 * ((z ** 2).sum(axis=2) + log_det_var)  # n_samples, k
    # nll = -np.log((np.exp(-log_p) * w).sum(axis=1))
    nll = -spe.logsumexp(-nlog_p, b=w, axis=1)
    return inds, x, nll
//  [python]
def classify(self, DTE):
        class_priors = self.class_priors
        
        log_densities0, log_densities1 = self.compute_lls(DTE)
        log_joint_densities0 = log_densities0 + numpy.log(class_priors[0])
        log_joint_densities1 = log_densities1 + numpy.log(class_priors[1])
        
        logScores = numpy.zeros((2, DTE.shape[1]))
        logScores[0, :] = log_joint_densities0
        logScores[1, :] = log_joint_densities1
    
        logMarginal = numpy.max(logScores, axis=0) + numpy.log(numpy.sum(numpy.exp(logScores - row(numpy.max(logScores, axis=0))), axis=0))
        logScores = logScores - row(logMarginal)
        scores = numpy.exp(logScores)
        predictedLabels = scores.argmax(axis=0)
        return predictedLabels
//  [python]
def classify(self, DTE):
        class_priors = self.class_priors
        
        log_densities0, log_densities1 = self.compute_lls(DTE)
        log_joint_densities0 = log_densities0 + numpy.log(class_priors[0])
        log_joint_densities1 = log_densities1 + numpy.log(class_priors[1])
        
        logScores = numpy.zeros((2, DTE.shape[1]))
        logScores[0, :] = log_joint_densities0
        logScores[1, :] = log_joint_densities1
    
        logMarginal = scipy.special.logsumexp(logScores, axis=0)
        logScores = logScores - row(logMarginal)
        scores = numpy.exp(logScores)
        predictedLabels = scores.argmax(axis=0)
        return predictedLabels
//  [python]
def log_likelihood(self, obs):
        log_lik = self.log_complete_likelihood(obs)
        return np.logaddexp.reduce(log_lik, axis=0)
//  [python]
def log_likelihood(self, obs):
        log_lik = self.log_complete_likelihood(obs)
        return logsumexp(log_lik, axis=0)
//  [python]
def test_segmented_cuda_logsumexp(precision, atol, rtol, seed, loc, shape):
    rng = np.random.default_rng(seed)

    values = rng.normal(size=shape, loc=loc) * 1000.0

    if precision == np.float32:
        summer = custom_ops.SegmentedSumExp_f32(shape[1], shape[0])
    else:
        summer = custom_ops.SegmentedSumExp_f64(shape[1], shape[0])

    test_vals = summer.logsumexp(values)
    assert len(test_vals) == shape[0]
    # Verify that the results are deterministic
    np.testing.assert_array_equal(test_vals, summer.logsumexp(values))

    for test_val, vals in zip(test_vals, values):
        np.testing.assert_allclose(s.logsumexp(vals), test_val, atol=atol, rtol=rtol)
//  [python]
def test_segmented_cuda_logsumexp(precision, atol, rtol, seed, loc, shape):
    rng = np.random.default_rng(seed)

    values = rng.normal(size=shape, loc=loc) * 1000.0

    if precision == np.float32:
        summer = custom_ops.SegmentedSumExp_f32(shape[1], shape[0])
    else:
        summer = custom_ops.SegmentedSumExp_f64(shape[1], shape[0])

    test_vals = summer.logsumexp(values)
    assert len(test_vals) == shape[0]
    # Verify that the results are deterministic
    np.testing.assert_array_equal(test_vals, summer.logsumexp(values))

    for test_val, vals in zip(test_vals, values):
        np.testing.assert_allclose(logsumexp(vals), test_val, atol=atol, rtol=rtol)
//  [python]
def responsibilities(self, obs):
        log_lik = self.log_complete_likelihood(obs)
        resp = np.exp(log_lik - logsumexp(log_lik, axis=0, keepdims=True))
        return resp
//  [python]
def expected_log_likelihood(self, obs):
        log_lik = self.expected_log_complete_likelihood(obs)
        return logsumexp(log_lik, axis=0)
//  [python]
def expected_responsibilities(self, obs):
        log_lik = self.expected_log_complete_likelihood(obs)
        resp = np.exp(log_lik - logsumexp(log_lik, axis=0, keepdims=True))
        return resp
//  [python]
def test_segmented_cuda_logsumexp_ragged_arrays(precision, atol, rtol, seed, loc, num_samples):
    rng = np.random.default_rng(seed)
    max_values_per_segment = 100
    values = [rng.normal(size=rng.integers(1, max_values_per_segment), loc=loc) * 1000.0 for _ in range(num_samples)]

    klass = custom_ops.SegmentedSumExp_f64
    if precision == np.float32:
        klass = custom_ops.SegmentedSumExp_f32

    summer = klass(max_values_per_segment, num_samples)

    test_vals = summer.logsumexp(values)
    assert len(test_vals) == num_samples
    # Verify that the results are deterministic
    np.testing.assert_array_equal(test_vals, summer.logsumexp(values))

    unsegmented_version = klass(max_values_per_segment, 1)
    for test_val, vals in zip(test_vals, values):
        ref_logsumexp = s.logsumexp(vals)
        np.testing.assert_allclose(ref_logsumexp, test_val, atol=atol, rtol=rtol)
        np.testing.assert_allclose(ref_logsumexp, unsegmented_version.logsumexp([vals])[0], atol=atol, rtol=rtol)
//  [python]
def test_segmented_cuda_logsumexp_ragged_arrays(precision, atol, rtol, seed, loc, num_samples):
    rng = np.random.default_rng(seed)
    max_values_per_segment = 100
    values = [rng.normal(size=rng.integers(1, max_values_per_segment), loc=loc) * 1000.0 for _ in range(num_samples)]

    klass = custom_ops.SegmentedSumExp_f64
    if precision == np.float32:
        klass = custom_ops.SegmentedSumExp_f32

    summer = klass(max_values_per_segment, num_samples)

    test_vals = summer.logsumexp(values)
    assert len(test_vals) == num_samples
    # Verify that the results are deterministic
    np.testing.assert_array_equal(test_vals, summer.logsumexp(values))

    unsegmented_version = klass(max_values_per_segment, 1)
    for test_val, vals in zip(test_vals, values):
        ref_logsumexp = logsumexp(vals)
        np.testing.assert_allclose(ref_logsumexp, test_val, atol=atol, rtol=rtol)
        np.testing.assert_allclose(ref_logsumexp, unsegmented_version.logsumexp([vals])[0], atol=atol, rtol=rtol)