jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def EOSConstraints2Prior(eos_path, out_path, Neos, Constraint):
    """
    Given a directory of EOSs and nmma.joint.Constraint,
    returns sorted EOSs and the corresponding prior

    Parameters
    ----------
    eos_path: str
        Path to the original EOSs
    out_path: str
        Path to the sorted EOSs
    Neos: int
        Number of EOSs
    Constraint: nmma.joint.Constraint
        The EOS constraint to be consider

    Returns
    -------
    prior: bilby.core.Prior
        Prior on the EOSs
    logNorm: float
        log normalization constant for the EOS prior
    """

    logLs = EOSConstraintsLoglikelihood(eos_path, Neos, Constraint)
    logNorm = np.log(np.sum(np.exp(logLs)))

    weights = np.exp(logLs - logNorm)
    sorted_weights = np.sort(weights)

    EOSSorting(eos_path, out_path, Neos, weights)

    xx = np.arange(0, Neos + 1)
    yy = np.concatenate((sorted_weights, sorted_weights[-1]))
    prior = Interped(xx, yy, minimum=0, maximum=Neos, name='EOS')

    return prior, logNorm
//  [python]
def EOSConstraints2Prior(eos_path, out_path, Neos, Constraint):
    """
    Given a directory of EOSs and nmma.joint.Constraint,
    returns sorted EOSs and the corresponding prior

    Parameters
    ----------
    eos_path: str
        Path to the original EOSs
    out_path: str
        Path to the sorted EOSs
    Neos: int
        Number of EOSs
    Constraint: nmma.joint.Constraint
        The EOS constraint to be consider

    Returns
    -------
    prior: bilby.core.Prior
        Prior on the EOSs
    logNorm: float
        log normalization constant for the EOS prior
    """

    logLs = EOSConstraintsLoglikelihood(eos_path, Neos, Constraint)
    logNorm = scipy.special.logsumexp(logLs)

    weights = np.exp(logLs - logNorm)
    sorted_weights = np.sort(weights)

    EOSSorting(eos_path, out_path, Neos, weights)

    xx = np.arange(0, Neos + 1)
    yy = np.concatenate((sorted_weights, sorted_weights[-1]))
    prior = Interped(xx, yy, minimum=0, maximum=Neos, name='EOS')

    return prior, logNorm
//  [python]
def expected_responsibilities(self, x, y):
        log_lik = self.expected_log_complete_likelihood(x, y)
        resp = np.exp(log_lik - logsumexp(log_lik, axis=0, keepdims=True))
        return resp
//  [python]
def meanfield_predictive_weights(self, x):
        log_cluster_gating = np.log(self.gating.posterior.mean())

        log_weights = np.zeros((self.cluster_size, self.mixture_size, len(x)))
        for m in range(self.cluster_size):
            log_weights[m] = np.expand_dims(np.log(self.components[m].gating.posterior.mean()), axis=1) \
                             + self.components[m].basis.log_posterior_predictive_gaussian(x) + log_cluster_gating[m]

        weights = np.exp(log_weights - logsumexp(log_weights, axis=(0, 1), keepdims=True))
        return weights
//  [python]
def meanfield_predictive_activation(self, x):
        x = np.reshape(x, (-1, self.input_dim))
        xx = self.input_transform.transform(x) if self.scale else x

        log_basis = np.zeros((self.cluster_size, self.mixture_size, len(xx)))
        for m in range(self.cluster_size):
            log_basis[m] = np.expand_dims(np.log(self.components[m].gating.posterior.mean()), axis=1)\
                           + self.components[m].basis.log_posterior_predictive_gaussian(xx)

        log_gating = np.log(self.gating.posterior.mean())

        log_activations = log_basis + np.expand_dims(log_gating, axis=(1, 2))
        activations = np.exp(log_activations - logsumexp(log_activations, axis=(0, 1), keepdims=True))
        return activations
//  [python]
def bayesian_log_likelihood_score(wrapped_model, x, y, **kwargs):
    x = x.astype(np.float32)
    y = y.astype(np.float32)

    scores = None
    nll = wrapped_model.model._get_neg_log_likelihood()
    posterior_draws = 1 if wrapped_model.model.map_mode else 50
    for _ in range(posterior_draws):
        np_nll = -nll(y, wrapped_model.model(x, training=False))
        res = np.expand_dims(np_nll, axis=0)
        scores = res if scores is None else np.concatenate([scores, res], axis=0)
        print(scores.shape)
    logsumexp = scipy.misc.logsumexp(scores, axis=0) - np.log(posterior_draws)
    print(logsumexp.shape)
    return logsumexp.mean()
//  [python]
def bayesian_log_likelihood_score(wrapped_model, x, y, **kwargs):
    x = x.astype(np.float32)
    y = y.astype(np.float32)

    scores = None
    nll = wrapped_model.model._get_neg_log_likelihood()
    posterior_draws = 1 if wrapped_model.model.map_mode else 50
    for _ in range(posterior_draws):
        np_nll = -nll(y, wrapped_model.model(x, training=False))
        res = np.expand_dims(np_nll, axis=0)
        scores = res if scores is None else np.concatenate([scores, res], axis=0)
        print(scores.shape)
    logsumexp = scipy.special.logsumexp(scores, axis=0) - np.log(posterior_draws)
    print(logsumexp.shape)
    return logsumexp.mean()
//  [python]
def analyze(labels, logits, num_bins, plot_save_path, temperature=1):

    '''
    Perform error calibration: ECE, MCE, reliability plot, etc.

    Args:
        labels = true labels
        logits = prediction logits
        num_bins = number of bins (M)
        plot_save_path = path to save reliability plot
        temperature = scale (T) for logit to prob conversion

    Returns:
        ece = expected calibration error
        mce = maximum calibration error
        nll = negative log-likelihood

    Effects:
        reliability plot is generated and saved
    '''
    
    # confidence-binning and compute bin-wise average accuracy and confidence
    bins, binned, bin_accs, bin_confs, bin_sizes = calc_bins(
        labels, logits, num_bins, temperature
    )

    # ece and mce
    ece = ( np.abs(bin_accs-bin_confs) * (bin_sizes/bin_sizes.sum()) ).sum()
    mce = np.abs(bin_accs-bin_confs).max()

    # nll score
    logits_ = logits/temperature
    nll = np.mean(
        np.log(np.sum(np.exp(logits_), axis=1)) - logits_[np.arange(len(labels)), labels]
    )

    # reliability plot
    draw_reliability_plot(
        bins, bin_accs, 
        plot_save_path, 
        title=f'Temperature = {temperature}', 
        ece=ece, mce=mce, nll=nll
    )

    return ece, mce, nll
//  [python]
def analyze(labels, logits, num_bins, plot_save_path, temperature=1):

    '''
    Perform error calibration: ECE, MCE, reliability plot, etc.

    Args:
        labels = true labels
        logits = prediction logits
        num_bins = number of bins (M)
        plot_save_path = path to save reliability plot
        temperature = scale (T) for logit to prob conversion

    Returns:
        ece = expected calibration error
        mce = maximum calibration error
        nll = negative log-likelihood

    Effects:
        reliability plot is generated and saved
    '''
    
    # confidence-binning and compute bin-wise average accuracy and confidence
    bins, binned, bin_accs, bin_confs, bin_sizes = calc_bins(
        labels, logits, num_bins, temperature
    )

    # ece and mce
    ece = ( np.abs(bin_accs-bin_confs) * (bin_sizes/bin_sizes.sum()) ).sum()
    mce = np.abs(bin_accs-bin_confs).max()

    # nll score
    logits_ = logits/temperature
    nll = np.mean(
        scipy.special.logsumexp(logits_, axis=1) - logits_[np.arange(len(labels)), labels]
    )

    # reliability plot
    draw_reliability_plot(
        bins, bin_accs, 
        plot_save_path, 
        title=f'Temperature = {temperature}', 
        ece=ece, mce=mce, nll=nll
    )

    return ece, mce, nll
//  [python]
def estimate_Z_lars(Z_batch_op, sess):
  """Batched estimate of Z for eval for LARS."""
  num_batches = int(FLAGS.lars_Z_num_eval_samples / FLAGS.lars_Z_batch_size)
  tf.logging.info("Evaluating Z with %d samples" %
                  FLAGS.lars_Z_num_eval_samples)
  log_Zs = np.empty([num_batches], dtype=np.float64)
  for i in range(num_batches):
    log_Zs[i] = sess.run(Z_batch_op)
    tf.logging.info("Batch %d/%d complete" % (i + 1, num_batches))

  log_Zs -= np.log(num_batches)
  log_Z_estimate = np.max(log_Zs) + np.log(np.sum(np.exp(log_Zs - np.max(log_Zs))))
  Z_estimate = np.exp(log_Z_estimate)
  tf.logging.info("Z estimate: %0.4f" % Z_estimate)
  return Z_estimate
//  [python]
def estimate_Z_lars(Z_batch_op, sess):
  """Batched estimate of Z for eval for LARS."""
  num_batches = int(FLAGS.lars_Z_num_eval_samples / FLAGS.lars_Z_batch_size)
  tf.logging.info("Evaluating Z with %d samples" %
                  FLAGS.lars_Z_num_eval_samples)
  log_Zs = np.empty([num_batches], dtype=np.float64)
  for i in range(num_batches):
    log_Zs[i] = sess.run(Z_batch_op)
    tf.logging.info("Batch %d/%d complete" % (i + 1, num_batches))

  log_Zs -= np.log(num_batches)
  log_Z_estimate = scipy.special.logsumexp(log_Zs)
  Z_estimate = np.exp(log_Z_estimate)
  tf.logging.info("Z estimate: %0.4f" % Z_estimate)
  return Z_estimate
//  [python]
def test_logsumexp_benchmark(size, axis, benchmark):
    X = pt.matrix("X")
    X_max = pt.max(X, axis=axis, keepdims=True)
    X_max = pt.switch(pt.isinf(X_max), 0, X_max)
    X_lse = pt.log(pt.sum(pt.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 = pytensor.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 = pt.matrix("X")
    X_max = pt.max(X, axis=axis, keepdims=True)
    X_max = pt.switch(pt.isinf(X_max), 0, X_max)
    X_lse = pt.log(pt.sum(pt.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 = pytensor.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 test_logsumexp(self):
        from scipy.special import logsumexp
        a = torch.randn(5, 4)
        a[0, 0] = inf
        a[1, :] = -inf
        actual = a.logsumexp(1)
        expected = logsumexp(a.numpy(), 1)
        self.assertEqual(expected.shape, actual.shape)
        self.assertTrue(np.allclose(expected, actual.numpy()))
        # check that out is actually inplace
        b = torch.zeros(5, 2)
        c = b[:, 0]
        torch.logsumexp(a, 1, out=c)
        self.assertTrue(np.allclose(expected, b[:, 0].numpy()))
//  [python]
def get_likelihoods(learned_param):

    with tf.Session() as sess:
        approx_model.load_param(learned_param, sess)
        
        removed_likelihoods_np = []
        remain_likelihoods_np = []

        for _ in range(n_iter):
            removed_likelihoods_np_i, remain_likelihoods_np_i = sess.run([removed_likelihoods, remain_likelihoods])

            removed_likelihoods_np_i = s.logsumexp(removed_likelihoods_np_i, axis=2, keepdims=True) - np.log(removed_likelihoods_np_i.shape[2])
            # (nclass, ndata, 1)

            remain_likelihoods_np_i = s.logsumexp(remain_likelihoods_np_i, axis=2, keepdims=True) - np.log(remain_likelihoods_np_i.shape[2])
            # (nclass, ndata, 1)

            removed_likelihoods_np.append(removed_likelihoods_np_i)
            remain_likelihoods_np.append(remain_likelihoods_np_i)
        
        removed_likelihoods_np = np.concatenate(removed_likelihoods_np, axis=2) # concat over the sample dimension
        remain_likelihoods_np = np.concatenate(remain_likelihoods_np, axis=2)
        
        removed_likelihoods_np = s.logsumexp(removed_likelihoods_np, axis=2) - np.log(removed_likelihoods_np.shape[2])
        remain_likelihoods_np = s.logsumexp(remain_likelihoods_np, axis=2) - np.log(remain_likelihoods_np.shape[2])

    return removed_likelihoods_np, remain_likelihoods_np
//  [python]
def get_likelihoods(learned_param):

    with tf.Session() as sess:
        approx_model.load_param(learned_param, sess)
        
        removed_likelihoods_np = []
        remain_likelihoods_np = []

        for _ in range(n_iter):
            removed_likelihoods_np_i, remain_likelihoods_np_i = sess.run([removed_likelihoods, remain_likelihoods])

            removed_likelihoods_np_i = scipy.special.logsumexp(removed_likelihoods_np_i, axis=2, keepdims=True) - np.log(removed_likelihoods_np_i.shape[2])
            # (nclass, ndata, 1)

            remain_likelihoods_np_i = scipy.special.logsumexp(remain_likelihoods_np_i, axis=2, keepdims=True) - np.log(remain_likelihoods_np_i.shape[2])
            # (nclass, ndata, 1)

            removed_likelihoods_np.append(removed_likelihoods_np_i)
            remain_likelihoods_np.append(remain_likelihoods_np_i)
        
        removed_likelihoods_np = np.concatenate(removed_likelihoods_np, axis=2) # concat over the sample dimension
        remain_likelihoods_np = np.concatenate(remain_likelihoods_np, axis=2)
        
        removed_likelihoods_np = scipy.special.logsumexp(removed_likelihoods_np, axis=2) - np.log(removed_likelihoods_np.shape[2])
        remain_likelihoods_np = scipy.special.logsumexp(remain_likelihoods_np, axis=2) - np.log(remain_likelihoods_np.shape[2])

    return removed_likelihoods_np, remain_likelihoods_np
//  [python]
def test_logsumexp_dim(self):
        from scipy.special import logsumexp
        self._test_dim_ops(
            lambda t, d: t.logsumexp(d),
            lambda n, d: logsumexp(n, d),
            use_integral=False)
//  [python]
def _log_likelihood(self, y, log_pdf):
        """
        Note: y shape is (..., D, N) and not (..., N, D) like in log_likelihood

        Args:
            y: Normalized observations with shape (..., D, N).
            log_pdf: shape (..., K, N)

        Returns:
            log_likelihood, scalar

        """
        *independent, channels, num_observations = y.shape

        # log_pdf.shape: *independent, speakers, num_observations

        # first: sum above the speakers
        # second: sum above time frequency in log domain
        log_likelihood = np.sum(np.log(np.sum(np.exp(log_pdf), axis=-2)))
        return log_likelihood
//  [python]
def _log_likelihood(self, y, log_pdf):
        """
        Note: y shape is (..., D, N) and not (..., N, D) like in log_likelihood

        Args:
            y: Normalized observations with shape (..., D, N).
            log_pdf: shape (..., K, N)

        Returns:
            log_likelihood, scalar

        """
        *independent, channels, num_observations = y.shape

        # log_pdf.shape: *independent, speakers, num_observations

        # first: sum above the speakers
        # second: sum above time frequency in log domain
        log_likelihood = np.sum(scipy.special.logsumexp(log_pdf, axis=-2))
        return log_likelihood
//  [python]
def logsumexp(da, dims=None, *, b=True, return_sign=False, **kwargs):
    """Wrap and extend :func:`scipy.special.logsumexp`.

    Usage examples available at :ref:`stats_tutorial`
    """
    if dims is None:
        dims = get_default_dims(da.dims)
    if not isinstance(dims, str):
        core_dims = dims
        axis = tuple(-1 - i for i in reversed(range(len(dims))))
    else:
        core_dims = [dims]
        axis = -1

    if return_sign:
        out_dims = [[], []]
    else:
        out_dims = [[]]

    b_dims = []
    b_dims_to_keep = []
    if isinstance(b, xr.DataArray):
        b_dims_to_keep = [d for d in b.dims if d in da.dims and d not in core_dims]
        b_dims = [d for d in b.dims if d in core_dims]
        out_dims[0].extend(b_dims_to_keep)

    try:
        return xr.apply_ufunc(
            lambda a, b, **kwargs: s.logsumexp(a, b=b, **kwargs),
            da,
            b,
            input_core_dims=[b_dims_to_keep + core_dims, b_dims_to_keep + b_dims],
            output_core_dims=out_dims,
            kwargs={"return_sign": return_sign, "axis": axis},
            **kwargs,
        )
    except ValueError:
        out_dims[0] = []
        return xr.apply_ufunc(
            lambda a, b, **kwargs: s.logsumexp(a, b=b, **kwargs),
            *xr.broadcast(da, b),
            input_core_dims=[core_dims, core_dims],
            output_core_dims=out_dims,
            kwargs={"return_sign": return_sign, "axis": axis},
            **kwargs,
        )
//  [python]
def logsumexp(da, dims=None, *, b=True, return_sign=False, **kwargs):
    """Wrap and extend :func:`scipy.special.logsumexp`.

    Usage examples available at :ref:`stats_tutorial`
    """
    if dims is None:
        dims = get_default_dims(da.dims)
    if not isinstance(dims, str):
        core_dims = dims
        axis = tuple(-1 - i for i in reversed(range(len(dims))))
    else:
        core_dims = [dims]
        axis = -1

    if return_sign:
        out_dims = [[], []]
    else:
        out_dims = [[]]

    b_dims = []
    b_dims_to_keep = []
    if isinstance(b, xr.DataArray):
        b_dims_to_keep = [d for d in b.dims if d in da.dims and d not in core_dims]
        b_dims = [d for d in b.dims if d in core_dims]
        out_dims[0].extend(b_dims_to_keep)

    try:
        return xr.apply_ufunc(
            lambda a, b, **kwargs: special.logsumexp(a, b=b, **kwargs),
            da,
            b,
            input_core_dims=[b_dims_to_keep + core_dims, b_dims_to_keep + b_dims],
            output_core_dims=out_dims,
            kwargs={"return_sign": return_sign, "axis": axis},
            **kwargs,
        )
    except ValueError:
        out_dims[0] = []
        return xr.apply_ufunc(
            lambda a, b, **kwargs: special.logsumexp(a, b=b, **kwargs),
            *xr.broadcast(da, b),
            input_core_dims=[core_dims, core_dims],
            output_core_dims=out_dims,
            kwargs={"return_sign": return_sign, "axis": axis},
            **kwargs,
        )
//  [python]
def testLogSumExpWhere(self, shape, dtype):
    rng = jtu.rand_default(self.rng())
    x = rng(shape, dtype)
    rng = jtu.rand_bool(self.rng())
    mask = rng(shape, bool)
    y_expected = osp_special.logsumexp(x[mask]) if mask.any() else -jnp.inf
    y_actual = lsp_special.logsumexp(x, where=mask)
    self.assertAllClose(y_expected, y_actual, check_dtypes=False)
//  [python]
def logsumexp(self, attrs = None):
        if attrs is None:
            return s.logsumexp(self.values)
        axes = self.domain.axes(attrs)
        values = s.logsumexp(self.values, axis=axes) 
        newdom = self.domain.marginalize(attrs)
        return Factor(newdom, values)
//  [python]
def logsumexp(self, attrs = None):
        if attrs is None:
            return logsumexp(self.values)
        axes = self.domain.axes(attrs)
        values = logsumexp(self.values, axis=axes) 
        newdom = self.domain.marginalize(attrs)
        return Factor(newdom, values)
//  [python]
def get_action_boltzman(qvals, alpha=0.01):
    batchsize, n_act = qvals.shape[1]
    exp_m = scipy.misc.logsumexp(qvals / alpha, axis=1)
    exp_m = np.exp(qvals / alpha - exp_m)
    actions = [np.random.choice(n_act, p=exp_m[i]) for i in range(batchsize)]
    return actions
//  [python]
def get_action_boltzman(qvals, alpha=0.01):
    batchsize, n_act = qvals.shape[1]
    exp_m = scipy.special.logsumexp(qvals / alpha, axis=1)
    exp_m = np.exp(qvals / alpha - exp_m)
    actions = [np.random.choice(n_act, p=exp_m[i]) for i in range(batchsize)]
    return actions
//  [python]
def aggregate_prediction_from_probabilities(probabilities,
                                            aggregation_method=METHOD_AVERAGE):
    """Determine the whole-image class prediction from patch probabilities.

  Args:
    probabilities: Numpy array of marginal probabilities, shape
     (batch_size, num_classes).
    aggregation_method: String, the method of aggregating the patch
      probabilities.

  Returns:
    A WholeImagePrediction object.

  Raises:
    ValueError: If the aggregation method is not valid.
  """
    certainties = certainties_from_probabilities(probabilities)

    certainty_dict = {
        'mean': numpy.round(numpy.mean(certainties), 3),
        'max': numpy.round(numpy.max(certainties), 3)
    }

    weights = certainties
    weights = None if numpy.sum(weights) == 0 else weights

    if aggregation_method == METHOD_AVERAGE:
        probabilities_aggregated = numpy.average(probabilities, 0, weights=weights)
    elif aggregation_method == METHOD_PRODUCT:
        # For i denoting index within batch and c the class:
        #   Q_c = product_over_i(p_c(i))
        # probabilities_aggregated = Q_c / sum_over_c(Q_c)
        # The following computes this using logs for numerical stability.
        sum_log_probabilities = numpy.sum(numpy.log(probabilities), 0)
        probabilities_aggregated = numpy.exp(
            sum_log_probabilities - numpy.max(sum_log_probabilities)) / numpy.sum(numpy.exp(sum_log_probabilities - numpy.max(sum_log_probabilities)))
    else:
        raise ValueError('Invalid aggregation method %s.' % aggregation_method)
    predicted_class = numpy.argmax(probabilities_aggregated)
    certainty_dict['aggregate'] = numpy.round(
        get_certainty(probabilities_aggregated), 3)
    certainty_dict['weighted'] = numpy.round(
        numpy.average(
            certainties, 0, weights=weights), 3)

    assert sorted(CERTAINTY_TYPES.values()) == sorted(certainty_dict.keys())

    return WholeImagePrediction(predicted_class, certainty_dict,
                                probabilities_aggregated)
//  [python]
def _complete_logprior(logprior):
  '''
  If you don't want to include a certain model in the posterior (e.g., perhaps
  you're computing pairwise probabilities for supervariants, and so don't want
  to allow them to cocluster unless they're the same supervariant), set the
  corresponding entry in prior to 0.
  '''
  if logprior is None:
    logprior = {}
    logtotal = -np.inf
  else:
    for K in logprior.keys():
      assert K in ALL_MODELS
      assert logprior[K] <= 0
    logtotal = s.logsumexp(list(logprior.values()))

  assert logtotal <= 0
  remaining = set(ALL_MODELS) - set(logprior.keys())
  if np.isclose(0, logtotal):
    remaining_val = -np.inf
  else:
    remaining_val = np.log(1 - np.exp(logtotal)) - np.log(len(remaining))
  for K in remaining:
      logprior[K] = remaining_val

  logprior_vals = -np.inf * np.ones(len(logprior))
  for K in logprior:
    logprior_vals[getattr(Models, K)] = logprior[K]
  assert np.isclose(0, s.logsumexp(logprior_vals))
  return logprior_vals
//  [python]
def _complete_logprior(logprior):
  '''
  If you don't want to include a certain model in the posterior (e.g., perhaps
  you're computing pairwise probabilities for supervariants, and so don't want
  to allow them to cocluster unless they're the same supervariant), set the
  corresponding entry in prior to 0.
  '''
  if logprior is None:
    logprior = {}
    logtotal = -np.inf
  else:
    for K in logprior.keys():
      assert K in ALL_MODELS
      assert logprior[K] <= 0
    logtotal = scipy.special.logsumexp(list(logprior.values()))

  assert logtotal <= 0
  remaining = set(ALL_MODELS) - set(logprior.keys())
  if np.isclose(0, logtotal):
    remaining_val = -np.inf
  else:
    remaining_val = np.log(1 - np.exp(logtotal)) - np.log(len(remaining))
  for K in remaining:
      logprior[K] = remaining_val

  logprior_vals = -np.inf * np.ones(len(logprior))
  for K in logprior:
    logprior_vals[getattr(Models, K)] = logprior[K]
  assert np.isclose(0, scipy.special.logsumexp(logprior_vals))
  return logprior_vals
//  [python]
def _gauss_kde_logpdf(x, samples, weights):
    """
    Utility function to compute the Gaussian log KDE probability at x from
    already whitened samples, possibly with weights.
    Normalization constants are ignored.
    """
    X = x-samples
    return scipy.special.logsumexp(-0.5*(X*X).sum(axis=1), b=weights)
//  [python]
def _gauss_ballkde_logpdf(x, samples, weights, distance_weights):
    """
    Utility function to compute the Gaussian log KDE probability
    with variable ball bandwidth at x from already whitened samples,
    possibly with weights. Each element has its own smoothing scale
    that is passed as `distance_weights`.
    Normalization constants are ignored.
    """
    X = x-samples
    return scipy.special.logsumexp(-0.5*(X*X).sum(axis=1)/distance_weights**2,
                                   b=weights)
//  [python]
def _gauss_ellkde_logpdf(x, samples, weights, distance_weights):
    """
    Utility function to compute the Gaussian log KDE probability
    with variable ellipsoid bandwidth at x from already whitened samples,
    possibly with weights. Each element has its own smoothing matrix
    that is passed as `distance_weights`.
    Normalization constants are ignored.
    """
    X = x-samples
    X = np.einsum('...j,...jk,...k', X, distance_weights, X)
    return scipy.special.logsumexp(-0.5*X, b=weights)
//  [python]
def testLogSumExp(self, shapes, dtype, axis,
                    keepdims, return_sign, use_b):
    if jnp.issubdtype(dtype, jnp.complexfloating) and scipy_version < (1, 13, 0):
      self.skipTest("logsumexp of complex input uses scipy 1.13.0 semantics.")
    if not jtu.test_device_matches(["cpu"]):
      rng = jtu.rand_some_inf_and_nan(self.rng())
    else:
      rng = jtu.rand_default(self.rng())
    # TODO(mattjj): test autodiff
    if use_b:
      def scipy_fun(array_to_reduce, scale_array):
        res = s.logsumexp(array_to_reduce, axis, keepdims=keepdims,
                                    return_sign=return_sign, b=scale_array)
        if dtype == np.int32:
          res = jax.tree.map(lambda x: x.astype('float32'), res)
        return res

      def lax_fun(array_to_reduce, scale_array):
        return lsp_special.logsumexp(array_to_reduce, axis, keepdims=keepdims,
                                     return_sign=return_sign, b=scale_array)

      args_maker = lambda: [rng(shapes[0], dtype), rng(shapes[1], dtype)]
    else:
      def scipy_fun(array_to_reduce):
        res = s.logsumexp(array_to_reduce, axis, keepdims=keepdims,
                                    return_sign=return_sign)
        if dtype == np.int32:
          res = jax.tree.map(lambda x: x.astype('float32'), res)
        return res

      def lax_fun(array_to_reduce):
        return lsp_special.logsumexp(array_to_reduce, axis, keepdims=keepdims,
                                     return_sign=return_sign)

      args_maker = lambda: [rng(shapes[0], dtype)]
    tol = (
        {np.float32: 2e-4, np.complex64: 2e-4}
        if jtu.test_device_matches(["tpu"])
        else None
    )
    self._CheckAgainstNumpy(scipy_fun, lax_fun, args_maker, rtol=tol, atol=tol)
    tol = {np.float32: 1E-6, np.float64: 1E-14}
    self._CompileAndCheck(lax_fun, args_maker, rtol=tol, atol=tol)
//  [python]
def testLogSumExp(self, shapes, dtype, axis,
                    keepdims, return_sign, use_b):
    if jnp.issubdtype(dtype, jnp.complexfloating) and scipy_version < (1, 13, 0):
      self.skipTest("logsumexp of complex input uses scipy 1.13.0 semantics.")
    if not jtu.test_device_matches(["cpu"]):
      rng = jtu.rand_some_inf_and_nan(self.rng())
    else:
      rng = jtu.rand_default(self.rng())
    # TODO(mattjj): test autodiff
    if use_b:
      def scipy_fun(array_to_reduce, scale_array):
        res = osp_special.logsumexp(array_to_reduce, axis, keepdims=keepdims,
                                    return_sign=return_sign, b=scale_array)
        if dtype == np.int32:
          res = jax.tree.map(lambda x: x.astype('float32'), res)
        return res

      def lax_fun(array_to_reduce, scale_array):
        return lsp_special.logsumexp(array_to_reduce, axis, keepdims=keepdims,
                                     return_sign=return_sign, b=scale_array)

      args_maker = lambda: [rng(shapes[0], dtype), rng(shapes[1], dtype)]
    else:
      def scipy_fun(array_to_reduce):
        res = osp_special.logsumexp(array_to_reduce, axis, keepdims=keepdims,
                                    return_sign=return_sign)
        if dtype == np.int32:
          res = jax.tree.map(lambda x: x.astype('float32'), res)
        return res

      def lax_fun(array_to_reduce):
        return lsp_special.logsumexp(array_to_reduce, axis, keepdims=keepdims,
                                     return_sign=return_sign)

      args_maker = lambda: [rng(shapes[0], dtype)]
    tol = (
        {np.float32: 2e-4, np.complex64: 2e-4}
        if jtu.test_device_matches(["tpu"])
        else None
    )
    self._CheckAgainstNumpy(scipy_fun, lax_fun, args_maker, rtol=tol, atol=tol)
    tol = {np.float32: 1E-6, np.float64: 1E-14}
    self._CompileAndCheck(lax_fun, args_maker, rtol=tol, atol=tol)
//  [python]
def partial_eval(self, X, n_samples=5):
    if n_samples < 1000:
      res, iwae = self.sess.run(
          (self.lHat, self.iwae),
          feed_dict={self.x: X, self.n_samples: n_samples})
      res = [iwae] + res
    else:  # special case to handle OOM
      assert n_samples % 100 == 0, "When using large # of samples, it must be divisble by 100"
      res = []
      for i in range(int(n_samples/100)):
        print('-'*100, '开始partial_eval_2')
        logF, = self.sess.run(
            (self.logF,),
            feed_dict={self.x: X, self.n_samples: 100})
        print('-'*100, '结束partial_eval_2')
        res.append(np.logaddexp.reduce(logF, axis=1))
      res = [np.mean(np.logaddexp.reduce(res, axis=0) - np.log(n_samples))]
    return res
//  [python]
def partial_eval(self, X, n_samples=5):
    if n_samples < 1000:
      res, iwae = self.sess.run(
          (self.lHat, self.iwae),
          feed_dict={self.x: X, self.n_samples: n_samples})
      res = [iwae] + res
    else:  # special case to handle OOM
      assert n_samples % 100 == 0, "When using large # of samples, it must be divisble by 100"
      res = []
      for i in range(int(n_samples/100)):
        print('-'*100, '开始partial_eval_2')
        logF, = self.sess.run(
            (self.logF,),
            feed_dict={self.x: X, self.n_samples: 100})
        print('-'*100, '结束partial_eval_2')
        res.append(logsumexp(logF, axis=1))
      res = [np.mean(logsumexp(res, axis=0) - np.log(n_samples))]
    return res
//  [python]
def _calc_val(self, trajectories):
        costs = trajectories["costs"].copy()
        actions = trajectories["actions"].copy()
        qvals = trajectories["qvals"].copy() if "qvals" in trajectories else None

        delta = actions - self.mean_action[None, :, :]
        control_costs = self._control_costs(delta)
        total_costs = costs + self.beta * control_costs        
        q_hat = self.calculate_returns(total_costs, qvals, self.gamma, self.td_lam)
        # traj_costs = cost_to_go(costs, self.gamma_seq)[:,0]
        # control_costs = self._control_costs(delta)
        # total_costs = traj_costs.copy() + self.beta * control_costs.copy()
        
		# calculate log-sum-exp
        # c = (-1.0/self.beta) * total_costs.copy()
        # cmax = np.max(c)
        # c -= cmax
        # c = np.exp(c)
        # val1 = cmax + np.log(np.sum(c)) - np.log(c.shape[0])
        # val1 = -self.beta * val1
        q_hat = q_hat[:,0]
        val = -self.beta * scipy.misc.logsumexp((-1.0/self.beta) * q_hat, b=(1.0/q_hat.shape[0]))
        return val
//  [python]
def _calc_val(self, trajectories):
        costs = trajectories["costs"].copy()
        actions = trajectories["actions"].copy()
        qvals = trajectories["qvals"].copy() if "qvals" in trajectories else None

        delta = actions - self.mean_action[None, :, :]
        control_costs = self._control_costs(delta)
        total_costs = costs + self.beta * control_costs        
        q_hat = self.calculate_returns(total_costs, qvals, self.gamma, self.td_lam)
        # traj_costs = cost_to_go(costs, self.gamma_seq)[:,0]
        # control_costs = self._control_costs(delta)
        # total_costs = traj_costs.copy() + self.beta * control_costs.copy()
        
		# calculate log-sum-exp
        # c = (-1.0/self.beta) * total_costs.copy()
        # cmax = np.max(c)
        # c -= cmax
        # c = np.exp(c)
        # val1 = cmax + np.log(np.sum(c)) - np.log(c.shape[0])
        # val1 = -self.beta * val1
        q_hat = q_hat[:,0]
        val = -self.beta * scipy.special.logsumexp((-1.0/self.beta) * q_hat, b=(1.0/q_hat.shape[0]))
        return val
//  [python]
def logsumexp(self, attrs=None):
        if attrs is None:
            return s.logsumexp(self.values)
        axes = self.domain.axes(attrs)
        values = s.logsumexp(self.values, axis=axes)
        newdom = self.domain.marginalize(attrs)
        return Factor(newdom, values)
//  [python]
def logsumexp(self, attrs=None):
        if attrs is None:
            return logsumexp(self.values)
        axes = self.domain.axes(attrs)
        values = logsumexp(self.values, axis=axes)
        newdom = self.domain.marginalize(attrs)
        return Factor(newdom, values)
//  [python]
def test_logsumexp2():
    array = np.log(np.arange(12)).reshape(3, 4)
    assert_equal(s.logsumexp(array, axis=-1), logsumexp2(array, axis=-1))
//  [python]
def test_logsumexp2():
    array = np.log(np.arange(12)).reshape(3, 4)
    assert_equal(logsumexp(array, axis=-1), logsumexp2(array, axis=-1))
//  [python]
def logpdf_GMM(X, gmm):
    S = numpy.zeros((len(gmm), X.shape[1]))
    
    for g in range(len(gmm)):
        #print (w,mu,C)
        (w, mu, C) = gmm[g]
        S[g, :] = logpdf_GAU_ND(X, mu, C) + numpy.log(w)
        
    logdens = scipy.misc.logsumexp(S, axis=0)
    
    return S, logdens
//  [python]
def logpdf_GMM(X, gmm):
    S = numpy.zeros((len(gmm), X.shape[1]))
    
    for g in range(len(gmm)):
        #print (w,mu,C)
        (w, mu, C) = gmm[g]
        S[g, :] = logpdf_GAU_ND(X, mu, C) + numpy.log(w)
        
    logdens = scipy.special.logsumexp(S, axis=0)
    
    return S, logdens
//  [python]
def log_bcjr_compute_llr(app, K, S):
    llr = numpy.zeros(K, dtype=numpy.float32)

    app = app.reshape((K, S, 2))
    llr = scipy.misc.logsumexp(app[:,:,0], axis=1) \
        - scipy.misc.logsumexp(app[:,:,1], axis=1)

    return llr
//  [python]
def log_bcjr_compute_llr(app, K, S):
    llr = numpy.zeros(K, dtype=numpy.float32)

    app = app.reshape((K, S, 2))
    llr = scipy.special.logsumexp(app[:,:,0], axis=1) \
        - scipy.special.logsumexp(app[:,:,1], axis=1)

    return llr
//  [python]
def lognormexp(values, dim=0):
    """Log of exponentiates and normalizes a Tensor/Variable/np.ndarray.

    input:
        values: Tensor/Variable/np.ndarray [dim_1, ..., dim_N]
        dim: n
    output:
        result: Tensor/Variable/np.ndarray [dim_1, ..., dim_N]
            where result[i_1, ..., i_N] =

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

    if isinstance(values, np.ndarray):
        log_denominator = np.log(np.sum(np.exp(values), axis=dim, keepdims=True))
        # log_numerator = values
        return values - log_denominator
    else:
        log_denominator = logsumexp(values, dim=dim, keepdim=True)
        # log_numerator = values
        return values - log_denominator.expand_as(values)
//  [python]
def optimize_eta(eta,mu,word_count,beta_doc): 
        """Optimizes the variational parameter eta given the likelihood and the gradient function
        """
        def f(eta, word_count, beta_doc):
            """objective for the variational update q(eta)
            We want to maximize f, but numpy only implements minimize, so we minimize -f
            """
            # precomputation
            eta_ = np.insert(eta, K-1, 0)
            Ndoc = int(np.sum(word_count))
            #formula 
            # from cpp implementation: 
            # log(expeta * betas) * doc_cts - ndoc * log(sum(expeta))
            part1 = np.dot(word_count,(eta_.max() + np.log(np.exp(eta_ - eta_.max())@beta_doc)))-Ndoc*np.log(np.sum(np.exp(eta_)))
            part2 = .5*(eta_[:-1]-mu).T@siginv@(eta_[:-1]-mu)
            return np.float32(part2 - part1)

        def df(eta, word_count, beta_doc):
            """gradient for the objective of the variational update q(etas)
            """
            # precomputation
            eta_ = np.insert(eta, K-1, 0)
            Ndoc = int(np.sum(word_count))
            theta = stable_softmax(eta_)
            phi = softmax_weights(eta_, beta_doc)
            #formula
            # part1 = np.delete(np.sum(phi * word_count,axis=1) - Ndoc*theta, K-1)
            part1 = np.delete(np.sum(phi * word_count,axis=1) - Ndoc*theta, K-1)
            part2 = siginv@(eta_[:-1]-mu)
            # We want to maximize f, but numpy only implements minimize, so we
            # minimize -f
            return np.float64(part2 - part1)
        return optimize.minimize(
            f, 
            x0=eta,
            args=(word_count, beta_doc),
            jac=df, 
            method="BFGS",
            options={'disp': True}
            )
//  [python]
def optimize_eta(eta,mu,word_count,beta_doc): 
        """Optimizes the variational parameter eta given the likelihood and the gradient function
        """
        def f(eta, word_count, beta_doc):
            """objective for the variational update q(eta)
            We want to maximize f, but numpy only implements minimize, so we minimize -f
            """
            # precomputation
            eta_ = np.insert(eta, K-1, 0)
            Ndoc = int(np.sum(word_count))
            #formula 
            # from cpp implementation: 
            # log(expeta * betas) * doc_cts - ndoc * log(sum(expeta))
            part1 = np.dot(word_count,(eta_.max() + np.log(np.exp(eta_ - eta_.max())@beta_doc)))-Ndoc*scipy.special.logsumexp(eta_)
            part2 = .5*(eta_[:-1]-mu).T@siginv@(eta_[:-1]-mu)
            return np.float32(part2 - part1)

        def df(eta, word_count, beta_doc):
            """gradient for the objective of the variational update q(etas)
            """
            # precomputation
            eta_ = np.insert(eta, K-1, 0)
            Ndoc = int(np.sum(word_count))
            theta = stable_softmax(eta_)
            phi = softmax_weights(eta_, beta_doc)
            #formula
            # part1 = np.delete(np.sum(phi * word_count,axis=1) - Ndoc*theta, K-1)
            part1 = np.delete(np.sum(phi * word_count,axis=1) - Ndoc*theta, K-1)
            part2 = siginv@(eta_[:-1]-mu)
            # We want to maximize f, but numpy only implements minimize, so we
            # minimize -f
            return np.float64(part2 - part1)
        return optimize.minimize(
            f, 
            x0=eta,
            args=(word_count, beta_doc),
            jac=df, 
            method="BFGS",
            options={'disp': True}
            )
//  [python]
def compute_rel_abun(x, denom):
    if x.ndim == 1:
        x = np.expand_dims(x, axis=0)
    z = np.hstack((x, np.zeros((x.shape[0], 1))))
    p  = np.exp(z - np.max(z, axis=1, keepdims=True) - np.log(np.sum(np.exp(z - np.max(z, axis=1, keepdims=True)), axis=1, keepdims=True)))
    for i in range(p.shape[1]-1, denom, -1):
        tmp = np.copy(p[:,i-1])
        p[:,i-1] = np.copy(p[:,i])
        p[:,i] = tmp
    return p
//  [python]
def compute_rel_abun(x, denom):
    if x.ndim == 1:
        x = np.expand_dims(x, axis=0)
    z = np.hstack((x, np.zeros((x.shape[0], 1))))
    p  = np.exp(z - logsumexp(z, axis=1, keepdims=True))
    for i in range(p.shape[1]-1, denom, -1):
        tmp = np.copy(p[:,i-1])
        p[:,i-1] = np.copy(p[:,i])
        p[:,i] = tmp
    return p
//  [python]
def func(observed_counts, counts, base_lambda, error_rate):
    sums = np.sum(counts, axis=-1)[:, None]
    frequencies = np.log(counts / sums)
    poisson_lambda = (np.arange(counts.shape[1])[None, :] + error_rate) * base_lambda
    prob = fast_poisson_logpmf(observed_counts[:, None], 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 func(observed_counts, counts, base_lambda, error_rate):
    sums = np.sum(counts, axis=-1)[:, None]
    frequencies = np.log(counts / sums)
    poisson_lambda = (np.arange(counts.shape[1])[None, :] + error_rate) * base_lambda
    prob = fast_poisson_logpmf(observed_counts[:, None], poisson_lambda)
    prob = logsumexp(frequencies + prob, axis=-1)
    return prob
//  [python]
def compute_nll_given_preacts(labels, preacts):
    log_sum_exp = np.log(np.sum(np.exp(preacts), axis=1))
    tsb_logits_trueclass = np.sum(preacts*labels, axis=1)
    log_likelihoods = tsb_logits_trueclass - log_sum_exp
    nll = -np.mean(log_likelihoods)
    return nll
//  [python]
def compute_nll_given_preacts(labels, preacts):
    log_sum_exp = scipy.special.logsumexp(a=preacts, axis=1) 
    tsb_logits_trueclass = np.sum(preacts*labels, axis=1)
    log_likelihoods = tsb_logits_trueclass - log_sum_exp
    nll = -np.mean(log_likelihoods)
    return nll
//  [python]
def GMM_EM_tiedDiag(X,gmm,psi= 0.01): #X -> ev
    '''
    EM algorithm for GMM tied diagonal covariance
    It estimates the parameters of a GMM that maximize the ll for
    a training set X
    If psi is given it's used for constraining the eigenvalues of the
    covariance matrices to be larger or equal to psi
    '''
    llNew = None
    llOld = None
    G = len(gmm)
    N = X.shape[1]
    sigma_array = []
    while llOld is None or llNew - llOld > 1e-6:
        llOld = llNew
        SJ = numpy.zeros((G, N))
        for g in range(G):
            SJ[g, :] = logpdf_GAU_ND(
                X, gmm[g][1], gmm[g][2]) + numpy.log(gmm[g][0])
        SM = numpy.max(SJ, axis=0) + numpy.log(numpy.sum(numpy.exp(SJ - numpy.max(SJ, axis=0)), axis=0))
        llNew = SM.sum()/N
        P = numpy.exp(SJ-SM)
        gmmNew = []
        sigmaTied=numpy.zeros((X.shape[0],X.shape[0]))
        for g in range(G):
            # m step
            gamma = P[g, :]
            Z = gamma.sum()
            F = (mrow(gamma)*X).sum(1)
            S = numpy.dot(X, (mrow(gamma)*X).T)
            w = Z/N
            mu = mcol(F/Z)
            #Sigma = S/Z - numpy.dot(mu, mu.T)
            #Sigma = Sigma * numpy.eye(Sigma.shape[0])
            #Sigma = Sigma * Z
            #gmmNew.append((w, mu, Sigma))
            sigma = S/Z - numpy.dot(mu, mu.T)
            sigmaTied += Z * sigma
            gmmNew.append((w, mu))  

        # calculate tied covariance
        gmm = gmmNew
        sigmaTied /= N
        sigmaTied *= numpy.eye(sigma.shape[0])
        U, s, _ = numpy.linalg.svd(sigmaTied)
        s[s<psi] = psi
        sigmaTied = numpy.dot(U, mcol(s)*U.T)
        
        newGmm = []
        for i in range(len(gmm)):
            (w, mu) = gmm[i]
            newGmm.append((w, mu, sigmaTied))
        
        gmm = newGmm
        #print(llNew,'llnew') #increase - if decrease problem
    #print(llNew-llOld,'llnew-llold')
    return gmm
//  [python]
def GMM_EM_tiedDiag(X,gmm,psi= 0.01): #X -> ev
    '''
    EM algorithm for GMM tied diagonal covariance
    It estimates the parameters of a GMM that maximize the ll for
    a training set X
    If psi is given it's used for constraining the eigenvalues of the
    covariance matrices to be larger or equal to psi
    '''
    llNew = None
    llOld = None
    G = len(gmm)
    N = X.shape[1]
    sigma_array = []
    while llOld is None or llNew - llOld > 1e-6:
        llOld = llNew
        SJ = numpy.zeros((G, N))
        for g in range(G):
            SJ[g, :] = logpdf_GAU_ND(
                X, gmm[g][1], gmm[g][2]) + numpy.log(gmm[g][0])
        SM = scipy.special.logsumexp(SJ, axis=0)
        llNew = SM.sum()/N
        P = numpy.exp(SJ-SM)
        gmmNew = []
        sigmaTied=numpy.zeros((X.shape[0],X.shape[0]))
        for g in range(G):
            # m step
            gamma = P[g, :]
            Z = gamma.sum()
            F = (mrow(gamma)*X).sum(1)
            S = numpy.dot(X, (mrow(gamma)*X).T)
            w = Z/N
            mu = mcol(F/Z)
            #Sigma = S/Z - numpy.dot(mu, mu.T)
            #Sigma = Sigma * numpy.eye(Sigma.shape[0])
            #Sigma = Sigma * Z
            #gmmNew.append((w, mu, Sigma))
            sigma = S/Z - numpy.dot(mu, mu.T)
            sigmaTied += Z * sigma
            gmmNew.append((w, mu))  

        # calculate tied covariance
        gmm = gmmNew
        sigmaTied /= N
        sigmaTied *= numpy.eye(sigma.shape[0])
        U, s, _ = numpy.linalg.svd(sigmaTied)
        s[s<psi] = psi
        sigmaTied = numpy.dot(U, mcol(s)*U.T)
        
        newGmm = []
        for i in range(len(gmm)):
            (w, mu) = gmm[i]
            newGmm.append((w, mu, sigmaTied))
        
        gmm = newGmm
        #print(llNew,'llnew') #increase - if decrease problem
    #print(llNew-llOld,'llnew-llold')
    return gmm
//  [python]
def f(eta, word_count, beta_doc_kv):
    # precomputation
    eta = np.insert(eta, K - 1, 0)
    Ndoc = int(np.sum(word_count))
    # formula
    # from cpp implementation:
    # log(expeta * betas) * doc_cts - ndoc * log(sum(expeta))
    return np.float64((0.5 * (eta[:-1] - mu).T @ siginv @ (eta[:-1] - mu)) - (np.dot(
        word_count, eta.max() + np.log(np.exp(eta - eta.max()) @ beta_doc_kv))
     - Ndoc * np.log(np.sum(np.exp(eta)))))
//  [python]
def f(eta, word_count, beta_doc_kv):
    # precomputation
    eta = np.insert(eta, K - 1, 0)
    Ndoc = int(np.sum(word_count))
    # formula
    # from cpp implementation:
    # log(expeta * betas) * doc_cts - ndoc * log(sum(expeta))
    return np.float64((0.5 * (eta[:-1] - mu).T @ siginv @ (eta[:-1] - mu)) - (np.dot(
        word_count, eta.max() + np.log(np.exp(eta - eta.max()) @ beta_doc_kv))
     - Ndoc * scipy.special.logsumexp(eta)))
//  [python]
def _get_optimal_ws_and_biases(self, preacts, labels):
         
        def eval_func(x):
            ws = np.array(x)

            vs_logits = preacts*ws[None,:]
            log_sum_exp = np.log(np.sum(np.exp(vs_logits), axis=1))
            exp_vs_logits = np.exp(vs_logits)
            sum_exp = np.sum(exp_vs_logits, axis=1)

            log_likelihoods = (np.sum(vs_logits*labels,axis=1)
                               - log_sum_exp)
            nll = -np.mean(log_likelihoods)

            grads_ws = preacts*(labels - (exp_vs_logits/sum_exp[:,None]))

            #multiply by -1 because we care about *negative* log likelihood
            mean_grads_ws = -np.mean(grads_ws, axis=0)

            return nll, mean_grads_ws

        if (self.verbose):
            original_nll = compute_nll(labels=labels, preacts=preacts,
                                       t=1.0, bs=np.zeros(labels.shape[1]))
            print("Original NLL is: ",original_nll)
            
        optimization_result = scipy.optimize.minimize(
                      fun=eval_func,
                      #fun=lambda x: eval_func(x)[0],
                      x0=np.array([1.0 for x in range(preacts.shape[1])]),
                      bounds=[(0,None) for x in range(preacts.shape[1])],
                      jac=True,
                      method='L-BFGS-B',
                      tol=1e-07,
                      **self.lbfgs_kwargs)
        if (self.verbose):
            print(optimization_result)
        assert optimization_result.success==True, optimization_result
        
        ws = optimization_result.x 
        return ws