jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def calculate_pac_dir(features_np_all, label_np_all, alpha=1.):
  """Compute the PACTran-Dirichlet estimator."""
  prob_np_all,_ = gmm_estimator(features_np_all, label_np_all)
#   starttime = time.time()
  label_np_all = one_hot(label_np_all)  # [n, v]
  soft_targets_sum = np.sum(label_np_all, axis=0)  # [v]
  soft_targets_sum = np.expand_dims(soft_targets_sum, axis=1)  # [v, 1]
  a0 = alpha * soft_targets_sum / np.sum(soft_targets_sum) + 1e-10

  # initialize
  qz = prob_np_all  # [n, d]
  log_s = np.log(prob_np_all + 1e-10)  # [n, d]

  for _ in range(10):
    aw = a0 + np.sum(np.einsum("BY,BZ->BYZ", label_np_all, qz), axis=0)
    logits_qz = (log_s +
                 np.matmul(label_np_all, scipy.special.digamma(aw)) -
                 np.reshape(scipy.special.digamma(np.sum(aw, axis=0)), [1, -1]))
    log_qz = logits_qz - np.max(logits_qz, axis=-1, keepdims=True)
    log_qz = log_qz - np.log(np.sum(np.exp(log_qz), axis=-1, keepdims=True))
    qz = np.exp(log_qz)

  log_c0 = scipy.special.loggamma(np.sum(a0)) - np.sum(
      scipy.special.loggamma(a0))
  log_c = scipy.special.loggamma(np.sum(aw, axis=0)) - np.sum(
      scipy.special.loggamma(aw), axis=0)

  pac_dir = np.sum(
      log_c0 - log_c - np.sum(qz * (log_qz - log_s), axis=0))
  pac_dir = -pac_dir / label_np_all.size
  return pac_dir
//  [python]
def calculate_pac_dir(features_np_all, label_np_all, alpha=1.):
  """Compute the PACTran-Dirichlet estimator."""
  prob_np_all,_ = gmm_estimator(features_np_all, label_np_all)
#   starttime = time.time()
  label_np_all = one_hot(label_np_all)  # [n, v]
  soft_targets_sum = np.sum(label_np_all, axis=0)  # [v]
  soft_targets_sum = np.expand_dims(soft_targets_sum, axis=1)  # [v, 1]
  a0 = alpha * soft_targets_sum / np.sum(soft_targets_sum) + 1e-10

  # initialize
  qz = prob_np_all  # [n, d]
  log_s = np.log(prob_np_all + 1e-10)  # [n, d]

  for _ in range(10):
    aw = a0 + np.sum(np.einsum("BY,BZ->BYZ", label_np_all, qz), axis=0)
    logits_qz = (log_s +
                 np.matmul(label_np_all, scipy.special.digamma(aw)) -
                 np.reshape(scipy.special.digamma(np.sum(aw, axis=0)), [1, -1]))
    log_qz = logits_qz - scipy.special.logsumexp(
        logits_qz, axis=-1, keepdims=True)
    qz = np.exp(log_qz)

  log_c0 = scipy.special.loggamma(np.sum(a0)) - np.sum(
      scipy.special.loggamma(a0))
  log_c = scipy.special.loggamma(np.sum(aw, axis=0)) - np.sum(
      scipy.special.loggamma(aw), axis=0)

  pac_dir = np.sum(
      log_c0 - log_c - np.sum(qz * (log_qz - log_s), axis=0))
  pac_dir = -pac_dir / label_np_all.size
  return pac_dir
//  [python]
def profile_logits_to_log_probs(logit_pred_profs, axis=2):
    """
    Converts the model's predicted profile logits into normalized probabilities
    via a softmax on the specified dimension (defaults to axis=2).
    Arguments:
        `logit_pred_profs`: a tensor/array containing the predicted profile
            logits
    Returns a tensor/array of the same shape, containing the predicted profiles
    as log probabilities by doing a log softmax on the specified dimension. If
    the input is a tensor, the output will be a tensor. If the input is a NumPy
    array, the output will be a NumPy array. Note that the  reason why this
    function returns log probabilities rather than raw probabilities is for
    numerical stability.
    """
    if type(logit_pred_profs) is np.ndarray:
        return logit_pred_profs - \
            np.log(np.sum(np.exp(logit_pred_profs), axis=axis, keepdims=True))
    else:
        return torch.log_softmax(logit_pred_profs, dim=axis)
//  [python]
def profile_logits_to_log_probs(logit_pred_profs, axis=2):
    """
    Converts the model's predicted profile logits into normalized probabilities
    via a softmax on the specified dimension (defaults to axis=2).
    Arguments:
        `logit_pred_profs`: a tensor/array containing the predicted profile
            logits
    Returns a tensor/array of the same shape, containing the predicted profiles
    as log probabilities by doing a log softmax on the specified dimension. If
    the input is a tensor, the output will be a tensor. If the input is a NumPy
    array, the output will be a NumPy array. Note that the  reason why this
    function returns log probabilities rather than raw probabilities is for
    numerical stability.
    """
    if type(logit_pred_profs) is np.ndarray:
        return logit_pred_profs - \
            scipy.special.logsumexp(logit_pred_profs, axis=axis, keepdims=True)
    else:
        return torch.log_softmax(logit_pred_profs, dim=axis)
//  [python]
def calculate_pac_gamma(features_np_all, label_np_all, alpha=1.):
    """Compute the PAC-Gamma estimator."""
    prob_np_all,_ = gmm_estimator(features_np_all, label_np_all)
    #   starttime = time.time()
    label_np_all = one_hot(label_np_all)  # [n, v]
    soft_targets_sum = np.sum(label_np_all, axis=0)  # [v]
    soft_targets_sum = np.expand_dims(soft_targets_sum, axis=1)  # [v, 1]

    a0 = alpha * soft_targets_sum / np.sum(soft_targets_sum) + 1e-10
    beta = 1.

    # initialize
    qz = prob_np_all  # [n, d]
    s = prob_np_all  # [n, d]
    log_s = np.log(prob_np_all + 1e-10)  # [n, d]
    aw = a0
    bw = beta
    lw = np.sum(s, axis=-1, keepdims=True) * np.sum(aw / bw)  # [n, 1]

    for _ in range(10):
        aw = a0 + np.sum(np.einsum("BY,BZ->BYZ", label_np_all, qz),
                        axis=0)  # [v, d]
        lw = np.matmul(
            s, np.expand_dims(np.sum(aw / bw, axis=0), axis=1))  # [n, 1]
        logits_qz = (
            log_s + np.matmul(label_np_all, scipy.special.digamma(aw) - np.log(bw)))
        log_qz = logits_qz - np.max(logits_qz, axis=-1, keepdims=True)
        log_qz = log_qz - np.log(np.sum(np.exp(log_qz), axis=-1, keepdims=True))
        qz = np.exp(log_qz)  # [n, a, d]

    pac_gamma = (
        np.sum(scipy.special.loggamma(a0) - scipy.special.loggamma(aw) +
                aw * np.log(bw) - a0 * np.log(beta)) +
        np.sum(np.sum(qz * (log_qz - log_s), axis=-1) +
                np.log(np.squeeze(lw, axis=-1)) - 1.))
    pac_gamma /= label_np_all.size
    pac_gamma += 1.
#   endtime = time.time()
    return pac_gamma
//  [python]
def calculate_pac_gamma(features_np_all, label_np_all, alpha=1.):
    """Compute the PAC-Gamma estimator."""
    prob_np_all,_ = gmm_estimator(features_np_all, label_np_all)
    #   starttime = time.time()
    label_np_all = one_hot(label_np_all)  # [n, v]
    soft_targets_sum = np.sum(label_np_all, axis=0)  # [v]
    soft_targets_sum = np.expand_dims(soft_targets_sum, axis=1)  # [v, 1]

    a0 = alpha * soft_targets_sum / np.sum(soft_targets_sum) + 1e-10
    beta = 1.

    # initialize
    qz = prob_np_all  # [n, d]
    s = prob_np_all  # [n, d]
    log_s = np.log(prob_np_all + 1e-10)  # [n, d]
    aw = a0
    bw = beta
    lw = np.sum(s, axis=-1, keepdims=True) * np.sum(aw / bw)  # [n, 1]

    for _ in range(10):
        aw = a0 + np.sum(np.einsum("BY,BZ->BYZ", label_np_all, qz),
                        axis=0)  # [v, d]
        lw = np.matmul(
            s, np.expand_dims(np.sum(aw / bw, axis=0), axis=1))  # [n, 1]
        logits_qz = (
            log_s + np.matmul(label_np_all, scipy.special.digamma(aw) - np.log(bw)))
        log_qz = logits_qz - scipy.special.logsumexp(
            logits_qz, axis=-1, keepdims=True)
        qz = np.exp(log_qz)  # [n, a, d]

    pac_gamma = (
        np.sum(scipy.special.loggamma(a0) - scipy.special.loggamma(aw) +
                aw * np.log(bw) - a0 * np.log(beta)) +
        np.sum(np.sum(qz * (log_qz - log_s), axis=-1) +
                np.log(np.squeeze(lw, axis=-1)) - 1.))
    pac_gamma /= label_np_all.size
    pac_gamma += 1.
#   endtime = time.time()
    return pac_gamma
//  [python]
def log_prob_parameters_mc(self, *, key, theta, n_samples=3e4):
        """Approximates p(theta) using Monte Carlo integration
            theta : parameters
        """

        logliks = []
        for tt in range(int(n_samples)):

            # sample from p(G)
            key, subk = random.split(key)
            g = self.g_dist.sample_G(key=subk)

            # evaluate log prob p(theta | G)
            logliks.append(self.log_prob_parameters(theta=theta, g=g))

            # print
            if not tt % int(n_samples / 1000) and tt > 0:
                curr = scipy.misc.logsumexp(
                    np.array(logliks[:tt + 1]) - np.log(tt + 1))
                print(f'iter = {tt}: log p(theta | G) [MC] = {curr}', end='\r')

        log_prob_obs = scipy.misc.logsumexp(
            np.array(logliks) - np.log(n_samples))
        return log_prob_obs
//  [python]
def log_prob_parameters_mc(self, *, key, theta, n_samples=3e4):
        """Approximates p(theta) using Monte Carlo integration
            theta : parameters
        """

        logliks = []
        for tt in range(int(n_samples)):

            # sample from p(G)
            key, subk = random.split(key)
            g = self.g_dist.sample_G(key=subk)

            # evaluate log prob p(theta | G)
            logliks.append(self.log_prob_parameters(theta=theta, g=g))

            # print
            if not tt % int(n_samples / 1000) and tt > 0:
                curr = scipy.special.logsumexp(
                    np.array(logliks[:tt + 1]) - np.log(tt + 1))
                print(f'iter = {tt}: log p(theta | G) [MC] = {curr}', end='\r')

        log_prob_obs = scipy.special.logsumexp(
            np.array(logliks) - np.log(n_samples))
        return log_prob_obs
//  [python]
def log_marginal_likelihood(self, *, x, all_g, z_g=None):
        """Computes log p(x) in closed form using conjugacy properties of Dirichlet-Categorical
            x : [n_samples, n_vars]
            all_g : list of all possible igraph.Graph objects in domain; is exhaustively summed over
        """

        # log p(x, G)
        log_prob_obs_g = np.zeros(len(all_g))

        # normalizing constant for log p(G) using exhaustive normalization
        if z_g is None:
            z_g = self.g_dist.log_normalization_constant(all_g=all_g)

        # log p(x, G)
        for i, g in enumerate(tqdm.tqdm(all_g, desc='p(X) log_marginal_likelihood', disable=not self.verbose)):

            # log p(x, G) = log (p(G)/Z) + log p(x | G)
            log_prob_obs_g[i] = self.g_dist.unnormalized_log_prob(g=g) - z_g \
                + self.log_marginal_likelihood_given_g(g=g, x=x)

        # log p(x) = log(sum_G exp(log p(x, G)))
        return np.logaddexp.reduce(log_prob_obs_g)
//  [python]
def log_marginal_likelihood(self, *, x, all_g, z_g=None):
        """Computes log p(x) in closed form using conjugacy properties of Dirichlet-Categorical
            x : [n_samples, n_vars]
            all_g : list of all possible igraph.Graph objects in domain; is exhaustively summed over
        """

        # log p(x, G)
        log_prob_obs_g = np.zeros(len(all_g))

        # normalizing constant for log p(G) using exhaustive normalization
        if z_g is None:
            z_g = self.g_dist.log_normalization_constant(all_g=all_g)

        # log p(x, G)
        for i, g in enumerate(tqdm.tqdm(all_g, desc='p(X) log_marginal_likelihood', disable=not self.verbose)):

            # log p(x, G) = log (p(G)/Z) + log p(x | G)
            log_prob_obs_g[i] = self.g_dist.unnormalized_log_prob(g=g) - z_g \
                + self.log_marginal_likelihood_given_g(g=g, x=x)

        # log p(x) = log(sum_G exp(log p(x, G)))
        return scipy.special.logsumexp(log_prob_obs_g)
//  [python]
def log_marginal_likelihood_given_g_mc(self, *, key, x, g, n_samples=3e4):
        """Approximates p(x | G) using Monte Carlo integration
            x : [n_samples, n_vars]
            g : graph
        """

        logliks = []
        for tt in range(int(n_samples)):

            # sample from p(theta | G)
            key, subk = random.split(key)
            theta = self.sample_parameters(key=subk, g=g)

            # evaluate likelihood log p(X | theta, G)
            logliks.append(self.log_likelihood(x=x, theta=theta, g=g))

            # print
            if not tt % int(n_samples / 1000) and tt > 0:
                curr = np.log(np.sum(np.exp(np.array(logliks[:tt + 1]) - np.log(tt + 1))))
                print(f'iter = {tt}: log p(X | G) [MC] = {curr}', end='\r')

        log_prob_obs = np.log(np.sum(np.exp(np.array(logliks) - np.log(n_samples))))
        return log_prob_obs
//  [python]
def log_marginal_likelihood_given_g_mc(self, *, key, x, g, n_samples=3e4):
        """Approximates p(x | G) using Monte Carlo integration
            x : [n_samples, n_vars]
            g : graph
        """

        logliks = []
        for tt in range(int(n_samples)):

            # sample from p(theta | G)
            key, subk = random.split(key)
            theta = self.sample_parameters(key=subk, g=g)

            # evaluate likelihood log p(X | theta, G)
            logliks.append(self.log_likelihood(x=x, theta=theta, g=g))

            # print
            if not tt % int(n_samples / 1000) and tt > 0:
                curr = scipy.special.logsumexp(
                    np.array(logliks[:tt + 1]) - np.log(tt + 1))
                print(f'iter = {tt}: log p(X | G) [MC] = {curr}', end='\r')

        log_prob_obs = scipy.special.logsumexp(
            np.array(logliks) - np.log(n_samples))
        return log_prob_obs
//  [python]
def log_marginal_likelihood_mc(self, *, key, x, n_samples=3e4):
        """Approximates normalization constant p(x) using Monte Carlo integration
            x : [n_samples, n_vars]
        """

        logliks = []
        for tt in range(int(n_samples)):

            # sample from p(G, theta) = p(G) p(theta | G)
            key, subk = random.split(key)
            g = self.g_dist.sample_G(key=subk)

            key, subk = random.split(key)
            theta = self.sample_parameters(key=subk, g=g)

            # evaluate likelihood log p(X | theta, G)
            logliks.append(self.log_likelihood(x=x, theta=theta, g=g))

            # print
            if not tt % int(n_samples / 1000) and tt > 0:
                curr = s.logsumexp(
                    (logliks[:tt + 1] - np.log(tt + 1)))
                print(f'iter = {tt}: log p(X) [MC] = {curr}', end='\r')
        print()
        log_prob_obs = s.logsumexp((logliks - np.log(n_samples)))
        return log_prob_obs
//  [python]
def log_marginal_likelihood_mc(self, *, key, x, n_samples=3e4):
        """Approximates normalization constant p(x) using Monte Carlo integration
            x : [n_samples, n_vars]
        """

        logliks = []
        for tt in range(int(n_samples)):

            # sample from p(G, theta) = p(G) p(theta | G)
            key, subk = random.split(key)
            g = self.g_dist.sample_G(key=subk)

            key, subk = random.split(key)
            theta = self.sample_parameters(key=subk, g=g)

            # evaluate likelihood log p(X | theta, G)
            logliks.append(self.log_likelihood(x=x, theta=theta, g=g))

            # print
            if not tt % int(n_samples / 1000) and tt > 0:
                curr = scipy.special.logsumexp(
                    (logliks[:tt + 1] - np.log(tt + 1)))
                print(f'iter = {tt}: log p(X) [MC] = {curr}', end='\r')
        print()
        log_prob_obs = scipy.special.logsumexp((logliks - np.log(n_samples)))
        return log_prob_obs
//  [python]
def _log_diff(log_p, log_q):
    # need to broadcast in case a or b is 0.5; logsumexp doesn't
    log_p, log_q = np.broadcast_arrays(log_p, log_q)
    return s.logsumexp([log_p, log_q+np.pi*1j], axis=0)
//  [python]
def _log_diff(log_p, log_q):
    # need to broadcast in case a or b is 0.5; logsumexp doesn't
    log_p, log_q = np.broadcast_arrays(log_p, log_q)
    return sc.logsumexp([log_p, log_q+np.pi*1j], axis=0)
//  [python]
def logsumexp(self, attr_set=None):
        if attr_set == None or len(attr_set) == 0:
            xp = cp.get_array_module(self.values)
            if xp == cp:
                values = xp.exp(self.values)
                values = xp.sum(values)
                values = xp.log(values)
                return values
            else:
                return s.logsumexp(self.values)
        assert(set(attr_set) <= set(self.domain.attr_list))
        sum_attr = list(set(self.domain.attr_list) - set(attr_set))
        sum_attr = tuple(self.domain.index_list(sum_attr))
        if cp.get_array_module(self.values) == cp:
            values = cp.exp(self.values)
            values = cp.sum(values, axis=sum_attr)
            values = cp.log(values)
        else:
            values = s.logsumexp(self.values, axis=sum_attr)
        return Factor(self.domain.project(attr_set), values, cp.get_array_module(self.values))
//  [python]
def logsumexp(self, attr_set=None):
        if attr_set == None or len(attr_set) == 0:
            xp = cp.get_array_module(self.values)
            if xp == cp:
                values = xp.exp(self.values)
                values = xp.sum(values)
                values = xp.log(values)
                return values
            else:
                return scipy.special.logsumexp(self.values)
        assert(set(attr_set) <= set(self.domain.attr_list))
        sum_attr = list(set(self.domain.attr_list) - set(attr_set))
        sum_attr = tuple(self.domain.index_list(sum_attr))
        if cp.get_array_module(self.values) == cp:
            values = cp.exp(self.values)
            values = cp.sum(values, axis=sum_attr)
            values = cp.log(values)
        else:
            values = scipy.special.logsumexp(self.values, axis=sum_attr)
        return Factor(self.domain.project(attr_set), values, cp.get_array_module(self.values))
//  [python]
def test_mixture_same_family_log_prob(self):
        probs = torch.rand(5, 5).softmax(dim=-1)
        loc = torch.randn(5, 5)
        scale = torch.rand(5, 5)

        def ref_log_prob(idx, x, log_prob):
            p = probs[idx].numpy()
            m = loc[idx].numpy()
            s = scale[idx].numpy()
            mix = scipy.stats.multinomial(1, p)
            comp = scipy.stats.norm(m, s)
            expected = np.logaddexp.reduce(comp.logpdf(x) + np.log(mix.p))
            self.assertEqual(log_prob, expected, atol=1e-3, rtol=0)

        self._check_log_prob(
            MixtureSameFamily(Categorical(probs=probs),
                              Normal(loc, scale)), ref_log_prob)
//  [python]
def test_mixture_same_family_log_prob(self):
        probs = torch.rand(5, 5).softmax(dim=-1)
        loc = torch.randn(5, 5)
        scale = torch.rand(5, 5)

        def ref_log_prob(idx, x, log_prob):
            p = probs[idx].numpy()
            m = loc[idx].numpy()
            s = scale[idx].numpy()
            mix = scipy.stats.multinomial(1, p)
            comp = scipy.stats.norm(m, s)
            expected = scipy.special.logsumexp(comp.logpdf(x) + np.log(mix.p))
            self.assertEqual(log_prob, expected, atol=1e-3, rtol=0)

        self._check_log_prob(
            MixtureSameFamily(Categorical(probs=probs),
                              Normal(loc, scale)), ref_log_prob)
//  [python]
def lognormexp(values, dim=0):
    """Exponentiates, normalizes and takes log of a torch.Tensor/np.ndarray.

    Args:
        values: torch.Tensor/np.ndarray [dim_1, ..., dim_N]
        dim: n
    Returns:
        result: torch.Tensor/np.ndarray [dim_1, ..., dim_N]
            where result[i_1, ..., i_N] =

                                 exp(values[i_1, ..., i_N])
            log( ------------------------------------------------------------ )
                    sum_{j = 1}^{dim_n} exp(values[i_1, ..., j, ..., i_N])
    """

    if isinstance(values, np.ndarray):
        log_denominator = np.max(values, axis=dim, keepdims=True) + np.log(np.sum(np.exp(values - np.max(values, axis=dim, keepdims=True)), axis=dim, keepdims=True))
        # log_numerator = values
        return values - log_denominator
    else:
        log_denominator = torch.logsumexp(values, dim=dim, keepdim=True)
        # log_numerator = values
        return values - log_denominator
//  [python]
def lognormexp(values, dim=0):
    """Exponentiates, normalizes and takes log of a torch.Tensor/np.ndarray.

    Args:
        values: torch.Tensor/np.ndarray [dim_1, ..., dim_N]
        dim: n
    Returns:
        result: torch.Tensor/np.ndarray [dim_1, ..., dim_N]
            where result[i_1, ..., i_N] =

                                 exp(values[i_1, ..., i_N])
            log( ------------------------------------------------------------ )
                    sum_{j = 1}^{dim_n} exp(values[i_1, ..., j, ..., i_N])
    """

    if isinstance(values, np.ndarray):
        log_denominator = scipy.special.logsumexp(
            values, axis=dim, keepdims=True
        )
        # log_numerator = values
        return values - log_denominator
    else:
        log_denominator = torch.logsumexp(values, dim=dim, keepdim=True)
        # log_numerator = values
        return values - log_denominator
//  [python]
def _test_logaddexp(self, device, dtype, base2):
        if base2:
            ref_func = np.logaddexp2
            our_func = torch.logaddexp2
        elif dtype in (torch.complex64, torch.complex128):
            # numpy has not implemented logaddexp for complex
            def _ref_func(x, y):
                return scipy.special.logsumexp(np.stack((x, y), axis=0), axis=0)

            ref_func = _ref_func
            our_func = torch.logaddexp
        else:
            ref_func = np.logaddexp
            our_func = torch.logaddexp

        def _test_helper(a, b):
            if dtype == torch.bfloat16:
                ref = ref_func(a.cpu().float().numpy(), b.cpu().float().numpy())
                v = our_func(a, b)
                self.assertEqual(ref, v.float(), atol=0.01, rtol=0.01)
            else:
                ref = ref_func(a.cpu().numpy(), b.cpu().numpy())
                v = our_func(a, b)
                self.assertEqual(ref, v)

        # simple test
        a = torch.randn(64, 2, dtype=dtype, device=device) - 0.5
        b = torch.randn(64, 2, dtype=dtype, device=device) - 0.5
        _test_helper(a, b)
        _test_helper(a[:3], b[:3])

        # large value test for numerical stability
        a *= 10000
        b *= 10000
        _test_helper(a, b)
        _test_helper(a[:3], b[:3])

        a = torch.tensor(
            [float("inf"), float("-inf"), float("inf"), float("nan")],
            dtype=dtype,
            device=device,
        )
        b = torch.tensor(
            [float("inf"), float("-inf"), float("-inf"), float("nan")],
            dtype=dtype,
            device=device,
        )
        _test_helper(a, b)
//  [python]
def calc_lh_grid(var1, var2):
  grid_step = 0.001
  N = 3*int(1/grid_step + 1) # N: number of grid points
  G = np.linspace(start=0, stop=1, num=N)[:,np.newaxis] # Nx1

  S = len(var1.total_reads) # S
  logprob_phi = generate_logprob_phi(N) # MxNxN
  logprob_models = np.nan * np.ones((S, NUM_MODELS)) # SxM

  for s in range(S):
    for modelidx in logprob_phi.keys():
      pv1, pv2 = [scipy.stats.binom.logpmf(V.var_reads[s], V.total_reads[s], V.omega_v[s]*G) for V in (var1, var2)] # Nx1
      p1 = np.tile(pv1, N)   # pv1 vector tiled as columns
      p2 = np.tile(pv2, N).T # pv2 vector tiled as rows
      P = p1 + p2 + logprob_phi[modelidx]
      logprob_models[s,modelidx] = np.max(P) + np.log(np.sum(np.exp(P - np.max(P))))

  return logprob_models
//  [python]
def calc_lh_grid(var1, var2):
  grid_step = 0.001
  N = 3*int(1/grid_step + 1) # N: number of grid points
  G = np.linspace(start=0, stop=1, num=N)[:,np.newaxis] # Nx1

  S = len(var1.total_reads) # S
  logprob_phi = generate_logprob_phi(N) # MxNxN
  logprob_models = np.nan * np.ones((S, NUM_MODELS)) # SxM

  for s in range(S):
    for modelidx in logprob_phi.keys():
      pv1, pv2 = [scipy.stats.binom.logpmf(V.var_reads[s], V.total_reads[s], V.omega_v[s]*G) for V in (var1, var2)] # Nx1
      p1 = np.tile(pv1, N)   # pv1 vector tiled as columns
      p2 = np.tile(pv2, N).T # pv2 vector tiled as rows
      P = p1 + p2 + logprob_phi[modelidx]
      logprob_models[s,modelidx] = scipy.special.logsumexp(P)

  return logprob_models
//  [python]
def averaged_rollout_metrics(experiment_folder: str):
    # Load folder
    rf_paths = [os.path.join(experiment_folder, r) for r in sorted(next(os.walk(experiment_folder))[1])]

    fail_nlls = np.stack([np.loadtxt(os.path.join(rfp, "fail_nlls.txt")) for rfp in rf_paths])
    fail_threshes = np.stack([np.loadtxt(os.path.join(rfp, "failure_threshes.txt")) for rfp in rf_paths])
    fails_per_stage = np.stack([np.loadtxt(os.path.join(rfp, "num_fails.txt")) for rfp in rf_paths])
    fail_probs = np.stack([np.loadtxt(os.path.join(rfp, "fail_probs.txt")) for rfp in rf_paths])

    fail_threshes_mu = fail_threshes.mean(0)
    fail_threshes_std = fail_threshes.std(0)
    # fails_per_stage_mu = fails_per_stage.mean(0)

    fail_nlls_mu = np.log(np.exp(fail_nlls).mean(0)) # Need special mean for logs
    fail_nlls_std = None

    fail_prob_mu = fail_probs.mean(0)
    fail_prob_sem = fail_probs.std(0) / fail_probs.shape[0]

    return fail_prob_mu, fail_prob_sem, fail_threshes_mu, fail_threshes_std, fail_nlls_mu, len(rf_paths)
//  [python]
def averaged_rollout_metrics(experiment_folder: str):
    # Load folder
    rf_paths = [os.path.join(experiment_folder, r) for r in sorted(next(os.walk(experiment_folder))[1])]

    fail_nlls = np.stack([np.loadtxt(os.path.join(rfp, "fail_nlls.txt")) for rfp in rf_paths])
    fail_threshes = np.stack([np.loadtxt(os.path.join(rfp, "failure_threshes.txt")) for rfp in rf_paths])
    fails_per_stage = np.stack([np.loadtxt(os.path.join(rfp, "num_fails.txt")) for rfp in rf_paths])
    fail_probs = np.stack([np.loadtxt(os.path.join(rfp, "fail_probs.txt")) for rfp in rf_paths])

    fail_threshes_mu = fail_threshes.mean(0)
    fail_threshes_std = fail_threshes.std(0)
    # fails_per_stage_mu = fails_per_stage.mean(0)

    fail_nlls_mu = scipy.special.logsumexp(fail_nlls, 0) - np.log(len(fail_nlls))  # Need special mean for logs
    fail_nlls_std = None

    fail_prob_mu = fail_probs.mean(0)
    fail_prob_sem = fail_probs.std(0) / fail_probs.shape[0]

    return fail_prob_mu, fail_prob_sem, fail_threshes_mu, fail_threshes_std, fail_nlls_mu, len(rf_paths)
//  [python]
def calc_lh_mc_2D(var1, var2):
  S = len(var1.total_reads) # S
  logprob_models = np.nan * np.ones((S, NUM_MODELS)) # SxM
  for s in range(S):
    for modelidx in (Models.cocluster, Models.A_B, Models.B_A, Models.diff_branches):
      mcsamps = int(1e6)
      phi1, phi2, area = _gen_samples(modelidx, S=mcsamps)

      logP = []
      for V, phi in ((var1, phi1), (var2, phi2)):
        logP.append(scipy.stats.binom.logpmf(V.var_reads[s], V.total_reads[s], V.omega_v[s]*phi))
      logprob_models[s,modelidx] = np.logaddexp.reduce(logP[0] + logP[1]) - np.log(mcsamps)

  return logprob_models
//  [python]
def calc_lh_mc_2D(var1, var2):
  S = len(var1.total_reads) # S
  logprob_models = np.nan * np.ones((S, NUM_MODELS)) # SxM
  for s in range(S):
    for modelidx in (Models.cocluster, Models.A_B, Models.B_A, Models.diff_branches):
      mcsamps = int(1e6)
      phi1, phi2, area = _gen_samples(modelidx, S=mcsamps)

      logP = []
      for V, phi in ((var1, phi1), (var2, phi2)):
        logP.append(scipy.stats.binom.logpmf(V.var_reads[s], V.total_reads[s], V.omega_v[s]*phi))
      logprob_models[s,modelidx] = scipy.special.logsumexp(logP[0] + logP[1]) - np.log(mcsamps)

  return logprob_models
//  [python]
def _calc_garbage_dumb(V1, V2):
  S = len(V1.total_reads) # S
  evidence = np.nan * np.ones(S)
  mcsamps = int(1e5)

  for sidx in range(S):
    logP = []
    for V in (V1, V2):
      phi = scipy.stats.uniform.rvs(loc=0, scale=1, size=mcsamps)
      B = scipy.stats.binom.logpmf(V.var_reads[sidx], V.total_reads[sidx], V.omega_v[sidx]*phi)
      logP.append(np.logaddexp.reduce(B) - np.log(mcsamps))
    evidence[sidx] = np.sum(logP)

  return evidence
//  [python]
def _calc_garbage_dumb(V1, V2):
  S = len(V1.total_reads) # S
  evidence = np.nan * np.ones(S)
  mcsamps = int(1e5)

  for sidx in range(S):
    logP = []
    for V in (V1, V2):
      phi = scipy.stats.uniform.rvs(loc=0, scale=1, size=mcsamps)
      B = scipy.stats.binom.logpmf(V.var_reads[sidx], V.total_reads[sidx], V.omega_v[sidx]*phi)
      logP.append(scipy.special.logsumexp(B) - np.log(mcsamps))
    evidence[sidx] = np.sum(logP)

  return evidence
//  [python]
def Z(self):
        return np.exp(scipy.misc.logsumexp(self.logits))
//  [python]
def Z(self):
        return np.exp(logsumexp(self.logits))
//  [python]
def get_logp_ub(inp, bsz=500, n_p=10, continuous=False, rce=False):
    elbos = []
    to_run = c_elbo if continuous else elbo
    if rce:
        to_run = to_run + tf.reduce_sum(kl, [1])
    for j in tqdm.trange(n_p):
        ee = []
        for k in range(0, inp.shape[0], bsz):
            ee.append(sess.run(to_run, {G.x_ph:inp[k:k+bsz], G.is_training_ph:False}))
        elbos.append(np.concatenate(ee, axis=0))
    logp = s.logsumexp(np.array(elbos), axis=0) - np.log(n_p)
    return logp / (inp.shape[1]**2 * inp.shape[3])
//  [python]
def get_logp_ub(inp, bsz=500, n_p=10, continuous=False, rce=False):
    elbos = []
    to_run = c_elbo if continuous else elbo
    if rce:
        to_run = to_run + tf.reduce_sum(kl, [1])
    for j in tqdm.trange(n_p):
        ee = []
        for k in range(0, inp.shape[0], bsz):
            ee.append(sess.run(to_run, {G.x_ph:inp[k:k+bsz], G.is_training_ph:False}))
        elbos.append(np.concatenate(ee, axis=0))
    logp = scipy.special.logsumexp(np.array(elbos), axis=0) - np.log(n_p)
    return logp / (inp.shape[1]**2 * inp.shape[3])
//  [python]
def calc_lh_mc_2D_dumb(var1, var2):
  S = len(var1.total_reads) # S
  logprob_models = np.nan * np.ones((S, NUM_MODELS)) # SxM
  for s in range(S):
    for modelidx in (Models.cocluster, Models.A_B, Models.B_A, Models.diff_branches):
      mcsamps = int(1e6)
      phi1 = scipy.stats.uniform.rvs(loc=0, scale=1, size=mcsamps)
      phi2 = scipy.stats.uniform.rvs(loc=0, scale=1, size=mcsamps)

      if modelidx == Models.A_B:
        prior = phi1 >= phi2
        area = 0.5
      elif modelidx == Models.B_A:
        prior = phi2 >= phi1
        area = 0.5
      elif modelidx == Models.diff_branches:
        prior = phi1 + phi2 <= 1
        area = 0.5
      elif modelidx == Models.cocluster:
        phi2 = phi1
        prior = np.array(mcsamps * [True])
        area = 1
      else:
        raise Exception('Unknown model: %s' % model)

      logP = []
      for V, phi in ((var1, phi1), (var2, phi2)):
        logP.append(scipy.stats.binom.logpmf(V.var_reads[s], V.total_reads[s], V.omega_v[s]*phi))
      logprob_models[s,modelidx] = np.logaddexp.reduce((logP[0] + logP[1] - np.log(area))[prior]) - np.log(mcsamps)

  return logprob_models
//  [python]
def calc_lh_mc_2D_dumb(var1, var2):
  S = len(var1.total_reads) # S
  logprob_models = np.nan * np.ones((S, NUM_MODELS)) # SxM
  for s in range(S):
    for modelidx in (Models.cocluster, Models.A_B, Models.B_A, Models.diff_branches):
      mcsamps = int(1e6)
      phi1 = scipy.stats.uniform.rvs(loc=0, scale=1, size=mcsamps)
      phi2 = scipy.stats.uniform.rvs(loc=0, scale=1, size=mcsamps)

      if modelidx == Models.A_B:
        prior = phi1 >= phi2
        area = 0.5
      elif modelidx == Models.B_A:
        prior = phi2 >= phi1
        area = 0.5
      elif modelidx == Models.diff_branches:
        prior = phi1 + phi2 <= 1
        area = 0.5
      elif modelidx == Models.cocluster:
        phi2 = phi1
        prior = np.array(mcsamps * [True])
        area = 1
      else:
        raise Exception('Unknown model: %s' % model)

      logP = []
      for V, phi in ((var1, phi1), (var2, phi2)):
        logP.append(scipy.stats.binom.logpmf(V.var_reads[s], V.total_reads[s], V.omega_v[s]*phi))
      logprob_models[s,modelidx] = scipy.special.logsumexp((logP[0] + logP[1] - np.log(area))[prior]) - np.log(mcsamps)

  return logprob_models
//  [python]
def main():
  import sys
  np.set_printoptions(linewidth=400, precision=3, threshold=sys.maxsize, suppress=True)
  np.seterr(divide='raise', invalid='raise')

  T = np.array([9000])
  x = np.arange(0, T[0] + 1)
  omega = np.array([0.5])
  phi = np.array([0])

  print('total_reads', T, sep='\t')
  print('omega', omega, sep='\t')
  print('phi', phi, sep='\t')
  limit = 20

  result_betabinom = logpx_betabinom(x, T, omega, phi)
  figs = []
  figs.append(px.line(x=x[:limit], y=np.exp(result_betabinom[:limit]), title='Beta binomial'))
  print(np.exp(s.logsumexp(result_betabinom)), np.exp(result_betabinom[:limit]))

  for epsilon in (1e-5, 1e-30):
    result = logpx_binom(x, T, omega, phi, epsilon=epsilon)
    fig = px.line(x=x[:limit], y=np.exp(result[:limit]), title=f'Binomial with epsilon={epsilon}')
    figs.append(fig)
    print(np.exp(s.logsumexp(result)), np.exp(result[:limit]))

  for F in figs:
    F.update_layout(yaxis = {'range': [0,1]})
//  [python]
def main():
  import sys
  np.set_printoptions(linewidth=400, precision=3, threshold=sys.maxsize, suppress=True)
  np.seterr(divide='raise', invalid='raise')

  T = np.array([9000])
  x = np.arange(0, T[0] + 1)
  omega = np.array([0.5])
  phi = np.array([0])

  print('total_reads', T, sep='\t')
  print('omega', omega, sep='\t')
  print('phi', phi, sep='\t')
  limit = 20

  result_betabinom = logpx_betabinom(x, T, omega, phi)
  figs = []
  figs.append(px.line(x=x[:limit], y=np.exp(result_betabinom[:limit]), title='Beta binomial'))
  print(np.exp(scipy.special.logsumexp(result_betabinom)), np.exp(result_betabinom[:limit]))

  for epsilon in (1e-5, 1e-30):
    result = logpx_binom(x, T, omega, phi, epsilon=epsilon)
    fig = px.line(x=x[:limit], y=np.exp(result[:limit]), title=f'Binomial with epsilon={epsilon}')
    figs.append(fig)
    print(np.exp(scipy.special.logsumexp(result)), np.exp(result[:limit]))

  for F in figs:
    F.update_layout(yaxis = {'range': [0,1]})
//  [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 smooth_svm_py(x, y, tau):
    x, y = to_numpy(x), to_numpy(y)
    n_samples, n_classes = x.shape
    scores = x + np.not_equal(np.arange(n_classes)[None, :], y[:, None]) - \
        x[np.arange(n_samples), y][:, None]
    loss = tau * np.mean(np.log(np.sum(np.exp(scores / tau), axis=1)))
    return loss
//  [python]
def smooth_svm_py(x, y, tau):
    x, y = to_numpy(x), to_numpy(y)
    n_samples, n_classes = x.shape
    scores = x + np.not_equal(np.arange(n_classes)[None, :], y[:, None]) - \
        x[np.arange(n_samples), y][:, None]
    loss = tau * np.mean(sp.logsumexp(scores / tau, axis=1))
    return loss
//  [python]
def conditional_sample(self, x, marginalize_idx, **kwargs):
        marginalization_backup = []
        component_posterior = np.zeros((self.num_components, x.shape[0]))
        for einet_counter, einet in enumerate(self.einets):
            marginalization_backup.append(einet.get_marginalization_idx())
            einet.set_marginalization_idx(marginalize_idx)
            lls = einet.forward(x)
            lls = lls.sum(1)
            component_posterior[einet_counter, :] = lls.detach().cpu().numpy() + np.log(self.p[einet_counter])

        component_posterior = component_posterior - logsumexp(component_posterior, 0, keepdims=True)
        component_posterior = np.exp(component_posterior)

        samples = np.zeros((x.shape[0], self.num_var, self.num_dims))
        for test_idx in range(x.shape[0]):
            component_idx = np.argmax(component_posterior[:, test_idx])
            sample = self.einets[component_idx].sample(x=x[test_idx:test_idx + 1, :], **kwargs)
            samples[test_idx, ...] = sample.squeeze().cpu().numpy()

        # restore the original marginalization indices
        for einet_counter, einet in enumerate(self.einets):
            einet.set_marginalization_idx(marginalization_backup[einet_counter])

        return samples
//  [python]
def __init__(self, n_vars, gibbs_temp=10., 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., 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 calc_lh_mc_1D(V1, V2):
  S = len(V1.total_reads) # S
  logprob_models = np.nan * np.ones((S, NUM_MODELS)) # SxM
  for sidx in range(S):
    for modelidx in (Models.cocluster, Models.A_B, Models.B_A, Models.diff_branches):
      mcsamps = int(1e5)
      phi1 = scipy.stats.uniform.rvs(loc=0, scale=1, size=mcsamps)

      if modelidx == Models.cocluster:
        logP = []
        for V in (V1, V2):
          logP.append(scipy.stats.binom.logpmf(V.var_reads[sidx], V.total_reads[sidx], V.omega_v[sidx]*phi1))
        logprob_models[sidx,modelidx] = np.logaddexp.reduce(logP[0] + logP[1]) - np.log(mcsamps)
      else:
        logP = scipy.stats.binom.logpmf(V1.var_reads[sidx], V1.total_reads[sidx], V1.omega_v[sidx]*phi1)

        lower = lhmath_native._make_lower(phi1, modelidx)
        upper = lhmath_native._make_upper(phi1, modelidx)
        A = V2.var_reads[sidx] + 1
        B = V2.ref_reads[sidx] + 1

        betainc = [scipy.special.betainc(A, B, V2.omega_v[sidx]*limit) for limit in (upper, lower)]
        logdenorm = scipy.special.betaln(A, B)
        # Add epsilon to ensure we don't take log of zero.
        # TODO: check if betainc[0] is close to betainc[1]. If so, the log will
        # likely produce -inf, so we can just set the LLH to -inf.
        logP += np.log(betainc[0] - betainc[1] + _EPSILON) + logdenorm

        logP = np.logaddexp.reduce(logP) - np.log(mcsamps)
        logP += np.log(2)
        logP += util.log_N_choose_K(V2.total_reads[sidx], V2.var_reads[sidx])
        logP -= np.log(V2.omega_v[sidx])
        logprob_models[sidx,modelidx] = logP

  return logprob_models
//  [python]
def calc_lh_mc_1D(V1, V2):
  S = len(V1.total_reads) # S
  logprob_models = np.nan * np.ones((S, NUM_MODELS)) # SxM
  for sidx in range(S):
    for modelidx in (Models.cocluster, Models.A_B, Models.B_A, Models.diff_branches):
      mcsamps = int(1e5)
      phi1 = scipy.stats.uniform.rvs(loc=0, scale=1, size=mcsamps)

      if modelidx == Models.cocluster:
        logP = []
        for V in (V1, V2):
          logP.append(scipy.stats.binom.logpmf(V.var_reads[sidx], V.total_reads[sidx], V.omega_v[sidx]*phi1))
        logprob_models[sidx,modelidx] = scipy.special.logsumexp(logP[0] + logP[1]) - np.log(mcsamps)
      else:
        logP = scipy.stats.binom.logpmf(V1.var_reads[sidx], V1.total_reads[sidx], V1.omega_v[sidx]*phi1)

        lower = lhmath_native._make_lower(phi1, modelidx)
        upper = lhmath_native._make_upper(phi1, modelidx)
        A = V2.var_reads[sidx] + 1
        B = V2.ref_reads[sidx] + 1

        betainc = [scipy.special.betainc(A, B, V2.omega_v[sidx]*limit) for limit in (upper, lower)]
        logdenorm = scipy.special.betaln(A, B)
        # Add epsilon to ensure we don't take log of zero.
        # TODO: check if betainc[0] is close to betainc[1]. If so, the log will
        # likely produce -inf, so we can just set the LLH to -inf.
        logP += np.log(betainc[0] - betainc[1] + _EPSILON) + logdenorm

        logP = scipy.special.logsumexp(logP) - np.log(mcsamps)
        logP += np.log(2)
        logP += util.log_N_choose_K(V2.total_reads[sidx], V2.var_reads[sidx])
        logP -= np.log(V2.omega_v[sidx])
        logprob_models[sidx,modelidx] = logP

  return logprob_models
//  [python]
def mce_partition_fh(
    env: resettable_env.TabularModelEnv,
    *,
    reward: Optional[np.ndarray] = None,
    discount: float = 1.0,
) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
    r"""Performs the soft Bellman backup for a finite-horizon MDP.

    Calculates V^{soft}, Q^{soft}, and pi using recurrences (9.1), (9.2), and
    (9.3) from Ziebart (2010).

    Args:
        env: a tabular, known-dynamics MDP.
        reward: a reward matrix. Defaults to env.reward_matrix.
        discount: discount rate.

    Returns:
        (V, Q, \pi) corresponding to the soft values, Q-values and MCE policy.
        V is a 2d array, indexed V[t,s]. Q is a 3d array, indexed Q[t,s,a].
        \pi is a 3d array, indexed \pi[t,s,a].
    """
    # shorthand
    horizon = env.horizon
    n_states = env.n_states
    n_actions = env.n_actions
    T = env.transition_matrix
    if reward is None:
        reward = env.reward_matrix

    # Initialization
    # indexed as V[t,s]
    V = np.full((horizon, n_states), -np.inf)
    # indexed as Q[t,s,a]
    Q = np.zeros((horizon, n_states, n_actions))
    broad_R = reward[:, None]

    # Base case: final timestep
    # final Q(s,a) is just reward
    Q[horizon - 1, :, :] = broad_R
    # V(s) is always normalising constant
    V[horizon - 1, :] = np.max(Q[horizon - 1, :, :], axis=1)

    # Recursive case
    for t in reversed(range(horizon - 1)):
        next_values_s_a = T @ V[t + 1, :]
        Q[t, :, :] = broad_R + discount * next_values_s_a
        V[t, :] = np.max(Q[t, :, :], axis=1)

    pi = np.exp(Q - V[:, :, None])

    return V, Q, pi
//  [python]
def mce_partition_fh(
    env: resettable_env.TabularModelEnv,
    *,
    reward: Optional[np.ndarray] = None,
    discount: float = 1.0,
) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
    r"""Performs the soft Bellman backup for a finite-horizon MDP.

    Calculates V^{soft}, Q^{soft}, and pi using recurrences (9.1), (9.2), and
    (9.3) from Ziebart (2010).

    Args:
        env: a tabular, known-dynamics MDP.
        reward: a reward matrix. Defaults to env.reward_matrix.
        discount: discount rate.

    Returns:
        (V, Q, \pi) corresponding to the soft values, Q-values and MCE policy.
        V is a 2d array, indexed V[t,s]. Q is a 3d array, indexed Q[t,s,a].
        \pi is a 3d array, indexed \pi[t,s,a].
    """
    # shorthand
    horizon = env.horizon
    n_states = env.n_states
    n_actions = env.n_actions
    T = env.transition_matrix
    if reward is None:
        reward = env.reward_matrix

    # Initialization
    # indexed as V[t,s]
    V = np.full((horizon, n_states), -np.inf)
    # indexed as Q[t,s,a]
    Q = np.zeros((horizon, n_states, n_actions))
    broad_R = reward[:, None]

    # Base case: final timestep
    # final Q(s,a) is just reward
    Q[horizon - 1, :, :] = broad_R
    # V(s) is always normalising constant
    V[horizon - 1, :] = scipy.special.logsumexp(Q[horizon - 1, :, :], axis=1)

    # Recursive case
    for t in reversed(range(horizon - 1)):
        next_values_s_a = T @ V[t + 1, :]
        Q[t, :, :] = broad_R + discount * next_values_s_a
        V[t, :] = scipy.special.logsumexp(Q[t, :, :], axis=1)

    pi = np.exp(Q - V[:, :, None])

    return V, Q, pi
//  [python]
def play(self) -> Tuple[Move, float]:
        assert not self.terminal
        # logsumexpを使ってオーバーフロー回避
        # value_n_exp = np.power(self.value_n, (1.0 / self.tree_config.play_temperature))
        # probs = value_n_exp / np.sum(value_n_exp)
        temp_log_value_n = (1.0 / self.tree_config.play_temperature) * np.log(self.value_n + 1e-20)
        denom = np.log(np.sum(np.exp(temp_log_value_n)))  # type: np.ndarray
        probs = np.exp(temp_log_value_n - denom)
        logger.info("Probs: {}".format([(self.move_list[i], probs[i]) for i in np.argsort(-probs)]))
        selected_edge = np.random.choice(np.arange(len(probs)), p=probs)
        return self.move_list[selected_edge], probs[selected_edge]
//  [python]
def play(self) -> Tuple[Move, float]:
        assert not self.terminal
        # logsumexpを使ってオーバーフロー回避
        # value_n_exp = np.power(self.value_n, (1.0 / self.tree_config.play_temperature))
        # probs = value_n_exp / np.sum(value_n_exp)
        temp_log_value_n = (1.0 / self.tree_config.play_temperature) * np.log(self.value_n + 1e-20)
        denom = scipy.special.logsumexp(temp_log_value_n)  # type: np.ndarray
        probs = np.exp(temp_log_value_n - denom)
        logger.info("Probs: {}".format([(self.move_list[i], probs[i]) for i in np.argsort(-probs)]))
        selected_edge = np.random.choice(np.arange(len(probs)), p=probs)
        return self.move_list[selected_edge], probs[selected_edge]
//  [python]
def _logpmf(observed_counts, counts, base_lambda, error_rate):
        #counts = counts.astype(np.float16)
        sums = np.sum(counts, axis=-1)[:, None]
        frequencies = np.log(counts / sums)
        poisson_lambda = (np.arange(counts.shape[1])[None,:] + error_rate) * base_lambda
        poisson_lambda = poisson_lambda.astype(np.float16)
        prob = fast_poisson_logpmf(observed_counts[:,None].astype(np.float16), poisson_lambda)
        prob = np.max(frequencies + prob, axis=-1) + np.log(np.sum(np.exp(frequencies + prob - np.max(frequencies + prob, axis=-1)[:,None]), axis=-1))
        return prob
//  [python]
def _logpmf(observed_counts, counts, base_lambda, error_rate):
        #counts = counts.astype(np.float16)
        sums = np.sum(counts, axis=-1)[:, None]
        frequencies = np.log(counts / sums)
        poisson_lambda = (np.arange(counts.shape[1])[None,:] + error_rate) * base_lambda
        poisson_lambda = poisson_lambda.astype(np.float16)
        prob = fast_poisson_logpmf(observed_counts[:,None].astype(np.float16), poisson_lambda)
        prob = logsumexp(frequencies + prob, axis=-1)
        return prob
//  [python]
def logmeanexp(x, axis=None, **kwargs):
    """
    Evaluate the logarithm of the mean of exponentiated values in a numerically stable
    fashion.

    Parameters
    ----------
    x : array_like
        Input array.
    axis : None or int or tuple[int], optional
        Axis or axes over which the mean is taken. By default `axis` is None, and all
        elements are averaged.
    **kwargs : dict
        Keyword arguments passed to `scipy.special.logsumexp`.

    Returns
    -------
    y : ndarray
        `np.log(np.mean(np.exp(a)))` calculated in a numerically more stable way.
    """
    lse = s.logsumexp(x, axis, **kwargs)
    if isinstance(axis, collections.abc.Iterable):
        size = np.prod([np.size(x, a) for a in axis])
    else:
        size = np.size(x, axis)
    return lse - np.log(size)