jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def first_derivative_log_nu_square(self, doc_lambda, doc_log_nu_square, doc_zeta_factor, total_word_count):
        assert doc_log_nu_square.shape == (self._number_of_topics,)
        assert doc_zeta_factor.shape == (self._number_of_topics, self._number_of_topics);

        exp_doc_log_nu_square = numpy.exp(doc_log_nu_square);

        exp_over_doc_zeta = scipy.special.logsumexp(
            doc_zeta_factor - doc_lambda[:, numpy.newaxis] - 0.5 * exp_doc_log_nu_square[:, numpy.newaxis], axis=1)
        exp_over_doc_zeta = numpy.exp(-exp_over_doc_zeta);

        if self._diagonal_covariance_matrix:
            first_derivative_log_nu_square = -0.5 * exp_doc_log_nu_square / self._alpha_sigma;
        else:
            first_derivative_log_nu_square = -0.5 * exp_doc_log_nu_square * numpy.diag(self._alpha_sigma_inv);
        first_derivative_log_nu_square += 0.5
        first_derivative_log_nu_square -= 0.5 * total_word_count * exp_over_doc_zeta * exp_doc_log_nu_square;

        return first_derivative_log_nu_square;
//  [python]
def _messages_backwards_log(trans_matrix,log_likelihoods):
        errs = np.seterr(over='ignore')
        Al = np.log(trans_matrix)
        aBl = log_likelihoods

        betal = np.zeros_like(aBl)

        for t in range(betal.shape[0]-2,-1,-1):
            betal[t] = logsumexp(Al + betal[t+1] + aBl[t+1],axis=1)

        np.seterr(**errs)
        return betal
//  [python]
def second_derivative_log_nu_square(self, doc_lambda, doc_log_nu_square, doc_zeta_factor, total_word_count):
        assert doc_log_nu_square.shape == (self._number_of_topics,)
        assert doc_zeta_factor.shape == (self._number_of_topics, self._number_of_topics);

        exp_doc_log_nu_square = numpy.exp(doc_log_nu_square);

        exp_over_doc_zeta = numpy.max(doc_zeta_factor - doc_lambda[:, numpy.newaxis] - 0.5 * exp_doc_log_nu_square[:, numpy.newaxis], axis=1) + numpy.log(numpy.sum(numpy.exp(doc_zeta_factor - doc_lambda[:, numpy.newaxis] - 0.5 * exp_doc_log_nu_square[:, numpy.newaxis] - exp_over_doc_zeta[:, numpy.newaxis]), axis=1));
        exp_over_doc_zeta = numpy.exp(-exp_over_doc_zeta);

        if self._diagonal_covariance_matrix:
            second_derivative_log_nu_square = -0.5 * exp_doc_log_nu_square / self._alpha_sigma;
        else:
            second_derivative_log_nu_square = -0.5 * exp_doc_log_nu_square * numpy.diag(self._alpha_sigma_inv);
        second_derivative_log_nu_square -= 0.5 * total_word_count * exp_over_doc_zeta * exp_doc_log_nu_square * (
                    1 + 0.5 * exp_doc_log_nu_square);

        return second_derivative_log_nu_square;
//  [python]
def second_derivative_log_nu_square(self, doc_lambda, doc_log_nu_square, doc_zeta_factor, total_word_count):
        assert doc_log_nu_square.shape == (self._number_of_topics,)
        assert doc_zeta_factor.shape == (self._number_of_topics, self._number_of_topics);

        exp_doc_log_nu_square = numpy.exp(doc_log_nu_square);

        exp_over_doc_zeta = scipy.special.logsumexp(
            doc_zeta_factor - doc_lambda[:, numpy.newaxis] - 0.5 * exp_doc_log_nu_square[:, numpy.newaxis], axis=1)
        exp_over_doc_zeta = numpy.exp(-exp_over_doc_zeta);

        if self._diagonal_covariance_matrix:
            second_derivative_log_nu_square = -0.5 * exp_doc_log_nu_square / self._alpha_sigma;
        else:
            second_derivative_log_nu_square = -0.5 * exp_doc_log_nu_square * numpy.diag(self._alpha_sigma_inv);
        second_derivative_log_nu_square -= 0.5 * total_word_count * exp_over_doc_zeta * exp_doc_log_nu_square * (
                    1 + 0.5 * exp_doc_log_nu_square);

        return second_derivative_log_nu_square;
//  [python]
def messages_backwards_log(self):
        betal = self._messages_backwards_log(self.trans_matrix,self.aBl)
        assert not np.isnan(betal).any()
        self._normalizer = np.logaddexp.reduce(np.log(self.pi_0) + betal[0] + self.aBl[0])
        return betal
//  [python]
def messages_backwards_log(self):
        betal = self._messages_backwards_log(self.trans_matrix,self.aBl)
        assert not np.isnan(betal).any()
        self._normalizer = logsumexp(np.log(self.pi_0) + betal[0] + self.aBl[0])
        return betal
//  [python]
def messages_forwards_log(self):
        alphal = self._messages_forwards_log(self.trans_matrix,self.pi_0,self.aBl)
        assert not np.any(np.isnan(alphal))
        self._normalizer = np.log(np.sum(np.exp(alphal[-1])))
        return alphal
//  [python]
def messages_forwards_log(self):
        alphal = self._messages_forwards_log(self.trans_matrix,self.pi_0,self.aBl)
        assert not np.any(np.isnan(alphal))
        self._normalizer = logsumexp(alphal[-1])
        return alphal
//  [python]
def _resample_r_from_mf(self):
        lognorm = np.log(np.sum(np.exp(self.rho_mf)))
        self.ridx = sample_discrete(np.exp(self.rho_mf - lognorm))
        self.r = self.r_support[self.ridx]
//  [python]
def _resample_r_from_mf(self):
        lognorm = logsumexp(self.rho_mf)
        self.ridx = sample_discrete(np.exp(self.rho_mf - lognorm))
        self.r = self.r_support[self.ridx]
//  [python]
def _messages_forwards_log(trans_matrix,init_state_distn,log_likelihoods):
        errs = np.seterr(over='ignore')
        Al = np.log(trans_matrix)
        aBl = log_likelihoods

        alphal = np.zeros_like(aBl)

        alphal[0] = np.log(init_state_distn) + aBl[0]
        for t in range(alphal.shape[0]-1):
            alphal[t+1] = np.max(alphal[t] + Al.T,axis=1) + np.log(np.sum(np.exp(alphal[t] + Al.T - np.max(alphal[t] + Al.T,axis=1,keepdims=True)),axis=1)) + aBl[t+1]

        np.seterr(**errs)
        return alphal
//  [python]
def _messages_forwards_log(trans_matrix,init_state_distn,log_likelihoods):
        errs = np.seterr(over='ignore')
        Al = np.log(trans_matrix)
        aBl = log_likelihoods

        alphal = np.zeros_like(aBl)

        alphal[0] = np.log(init_state_distn) + aBl[0]
        for t in range(alphal.shape[0]-1):
            alphal[t+1] = logsumexp(alphal[t] + Al.T,axis=1) + aBl[t+1]

        np.seterr(**errs)
        return alphal
//  [python]
def expected_log_likelihood(self,x):
        lognorm = s.logsumexp(self.rho_mf)
        return sum(np.exp(rho-lognorm)*d.expected_log_likelihood(x)
                for rho,d in zip(self.rho_mf,self._fixedr_distns))
//  [python]
def expected_log_likelihood(self,x):
        lognorm = logsumexp(self.rho_mf)
        return sum(np.exp(rho-lognorm)*d.expected_log_likelihood(x)
                for rho,d in zip(self.rho_mf,self._fixedr_distns))
//  [python]
def expcal(du, kbt):
    # ln<exp(X)> = ln (1/N sum[exp(x)]) = logsumexp(x) - log(N)
    lnexp = scipy.misc.logsumexp(du) - numpy.log(len(du))
    return - 1./kbt * lnexp
//  [python]
def expcal(du, kbt):
    # ln<exp(X)> = ln (1/N sum[exp(x)]) = logsumexp(x) - log(N)
    lnexp = scipy.special.logsumexp(du) - numpy.log(len(du))
    return - 1./kbt * lnexp
//  [python]
def _expected_statistics_from_messages(trans_potential,likelihood_log_potential,alphal,betal):
        expected_states = alphal + betal
        expected_states -= expected_states.max(1)[:,na]
        np.exp(expected_states,out=expected_states)
        expected_states /= expected_states.sum(1)[:,na]

        Al = np.log(trans_potential)
        log_joints = alphal[:-1,:,na] + (betal[1:,na,:] + likelihood_log_potential[1:,na,:]) + Al[na,...]
        log_joints -= log_joints.max((1,2))[:,na,na]
        joints = np.exp(log_joints)
        joints /= joints.sum((1,2))[:,na,na] # NOTE: renormalizing each isnt really necessary
        expected_transcounts = joints.sum(0)

        normalizer = s.logsumexp(alphal[0] + betal[0])

        return expected_states, expected_transcounts, normalizer
//  [python]
def _expected_statistics_from_messages(trans_potential,likelihood_log_potential,alphal,betal):
        expected_states = alphal + betal
        expected_states -= expected_states.max(1)[:,na]
        np.exp(expected_states,out=expected_states)
        expected_states /= expected_states.sum(1)[:,na]

        Al = np.log(trans_potential)
        log_joints = alphal[:-1,:,na] + (betal[1:,na,:] + likelihood_log_potential[1:,na,:]) + Al[na,...]
        log_joints -= log_joints.max((1,2))[:,na,na]
        joints = np.exp(log_joints)
        joints /= joints.sum((1,2))[:,na,na] # NOTE: renormalizing each isnt really necessary
        expected_transcounts = joints.sum(0)

        normalizer = logsumexp(alphal[0] + betal[0])

        return expected_states, expected_transcounts, normalizer
//  [python]
def hessian_direction_approximation_log_nu_square(self, doc_lambda, doc_log_nu_square, doc_zeta_factor,
                                                      total_word_count, direction_vector, epsilon=1e-6):
        assert doc_lambda.shape == (self._number_of_topics,)
        assert doc_log_nu_square.shape == (self._number_of_topics,)
        assert doc_zeta_factor.shape == (self._number_of_topics, self._number_of_topics)
        assert direction_vector.shape == (self._number_of_topics,)

        exp_doc_log_nu_square = numpy.exp(doc_log_nu_square);
        exp_doc_log_nu_square_epsilon_direction = numpy.exp(doc_log_nu_square + direction_vector * epsilon);

        log_exp_over_doc_zeta_epsilon_direction = numpy.max(
            doc_zeta_factor - doc_lambda[:, numpy.newaxis] - 0.5 * exp_doc_log_nu_square_epsilon_direction[:,
                                                                   numpy.newaxis], axis=1) + numpy.log(numpy.sum(
            numpy.exp(doc_zeta_factor - doc_lambda[:, numpy.newaxis] - 0.5 * exp_doc_log_nu_square_epsilon_direction[:,
                                                                   numpy.newaxis] - log_exp_over_doc_zeta_epsilon_direction[:, numpy.newaxis]), axis=1))
        log_exp_over_doc_zeta = numpy.max(
            doc_zeta_factor - doc_lambda[:, numpy.newaxis] - 0.5 * exp_doc_log_nu_square[:, numpy.newaxis], axis=1) + numpy.log(numpy.sum(
            numpy.exp(doc_zeta_factor - doc_lambda[:, numpy.newaxis] - 0.5 * exp_doc_log_nu_square[:, numpy.newaxis] - log_exp_over_doc_zeta[:, numpy.newaxis]), axis=1))

        # hessian_direction_log_nu_square = 0.5 * total_word_count * numpy.exp(log_doc_nu_square - log_exp_over_doc_zeta)
        # hessian_direction_log_nu_square += - 0.5 * total_word_count * numpy.exp(log_doc_nu_square + direction_vector * epsilon - log_exp_over_doc_zeta_epsilon_direction)

        hessian_direction_log_nu_square = 1 - numpy.exp(
            direction_vector * epsilon - log_exp_over_doc_zeta_epsilon_direction + log_exp_over_doc_zeta)
        hessian_direction_log_nu_square *= 0.5 * total_word_count * numpy.exp(doc_log_nu_square - log_exp_over_doc_zeta)

        if self._diagonal_covariance_matrix:
            hessian_direction_log_nu_square += 0.5 * (
                        exp_doc_log_nu_square - exp_doc_log_nu_square_epsilon_direction) / self._alpha_sigma;
        else:
            hessian_direction_log_nu_square += 0.5 * (
                        exp_doc_log_nu_square - exp_doc_log_nu_square_epsilon_direction) * numpy.diag(
                self._alpha_sigma_inv);

        hessian_direction_log_nu_square /= epsilon

        return hessian_direction_log_nu_square
//  [python]
def hessian_direction_approximation_log_nu_square(self, doc_lambda, doc_log_nu_square, doc_zeta_factor,
                                                      total_word_count, direction_vector, epsilon=1e-6):
        assert doc_lambda.shape == (self._number_of_topics,)
        assert doc_log_nu_square.shape == (self._number_of_topics,)
        assert doc_zeta_factor.shape == (self._number_of_topics, self._number_of_topics)
        assert direction_vector.shape == (self._number_of_topics,)

        exp_doc_log_nu_square = numpy.exp(doc_log_nu_square);
        exp_doc_log_nu_square_epsilon_direction = numpy.exp(doc_log_nu_square + direction_vector * epsilon);

        log_exp_over_doc_zeta_epsilon_direction = scipy.special.logsumexp(
            doc_zeta_factor - doc_lambda[:, numpy.newaxis] - 0.5 * exp_doc_log_nu_square_epsilon_direction[:,
                                                                   numpy.newaxis], axis=1)
        log_exp_over_doc_zeta = scipy.special.logsumexp(
            doc_zeta_factor - doc_lambda[:, numpy.newaxis] - 0.5 * exp_doc_log_nu_square[:, numpy.newaxis], axis=1)

        # hessian_direction_log_nu_square = 0.5 * total_word_count * numpy.exp(log_doc_nu_square - log_exp_over_doc_zeta)
        # hessian_direction_log_nu_square += - 0.5 * total_word_count * numpy.exp(log_doc_nu_square + direction_vector * epsilon - log_exp_over_doc_zeta_epsilon_direction)

        hessian_direction_log_nu_square = 1 - numpy.exp(
            direction_vector * epsilon - log_exp_over_doc_zeta_epsilon_direction + log_exp_over_doc_zeta)
        hessian_direction_log_nu_square *= 0.5 * total_word_count * numpy.exp(doc_log_nu_square - log_exp_over_doc_zeta)

        if self._diagonal_covariance_matrix:
            hessian_direction_log_nu_square += 0.5 * (
                        exp_doc_log_nu_square - exp_doc_log_nu_square_epsilon_direction) / self._alpha_sigma;
        else:
            hessian_direction_log_nu_square += 0.5 * (
                        exp_doc_log_nu_square - exp_doc_log_nu_square_epsilon_direction) * numpy.diag(
                self._alpha_sigma_inv);

        hessian_direction_log_nu_square /= epsilon

        return hessian_direction_log_nu_square
//  [python]
def pop_charge(pH, formal_charge, state_details):
    free_energies = []
    for item in state_details: #state_details [('SM42_micro001', 0.5304000000000001, -1, 0.0, 1.3872000000000002)
        if item[2] == formal_charge:
            free_energies.append(-beta * DeltaG(pH, item[0], state_details))
    if formal_charge == 0:
        free_energies.append(0 * pH)
    #print("free_energies",free_energies)
    return np.exp(s.logsumexp(free_energies))
//  [python]
def pop_charge(pH, formal_charge, state_details):
    free_energies = []
    for item in state_details: #state_details [('SM42_micro001', 0.5304000000000001, -1, 0.0, 1.3872000000000002)
        if item[2] == formal_charge:
            free_energies.append(-beta * DeltaG(pH, item[0], state_details))
    if formal_charge == 0:
        free_energies.append(0 * pH)
    #print("free_energies",free_energies)
    return np.exp(logsumexp(free_energies))
//  [python]
def logz_true(self):
        logz = scipy.misc.logsumexp(
            -self.alpha**(2 * np.arange(100) / self.ndim) / (8 * self.sig**2) +
            np.arange(100) * np.log(self.alpha) + np.log((1 - self.alpha)))
        return logz
//  [python]
def logz_true(self):
        logz = scipy.special.logsumexp(
            -self.alpha**(2 * np.arange(100) / self.ndim) / (8 * self.sig**2) +
            np.arange(100) * np.log(self.alpha) + np.log((1 - self.alpha)))
        return logz
//  [python]
def _accumulate_statistics(self, X: Dict[str, np.ndarray]) -> float:
        """Acccumulate the counts based on the sufficient statistics"""
        
        # Compute its current log-likelihood
        framelogprob = self._get_log_likelihood(X)
        
        # We add the the prior state log-probabilities
        posteriors = np.exp(framelogprob + np.log(self.prior_probs.T)) #type:ignore
        
        # Update the start counts
        self.state_counts += posteriors.sum(axis=0) #type:ignore

        # Updating the emission counts
        for src in X:
            self.emit_counts[src] += np.dot(posteriors.T, X[src]) #type:ignore

        for src, src2 in self.corr_counts:
            if src in X and src2 in X:
                self.corr_counts[(src, src2)] += np.dot(X[src2].T, X[src])
                
        return np.log(np.sum(np.exp(posteriors),axis=1)).sum()
//  [python]
def _accumulate_statistics(self, X: Dict[str, np.ndarray]) -> float:
        """Acccumulate the counts based on the sufficient statistics"""
        
        # Compute its current log-likelihood
        framelogprob = self._get_log_likelihood(X)
        
        # We add the the prior state log-probabilities
        posteriors = np.exp(framelogprob + np.log(self.prior_probs.T)) #type:ignore
        
        # Update the start counts
        self.state_counts += posteriors.sum(axis=0) #type:ignore

        # Updating the emission counts
        for src in X:
            self.emit_counts[src] += np.dot(posteriors.T, X[src]) #type:ignore

        for src, src2 in self.corr_counts:
            if src in X and src2 in X:
                self.corr_counts[(src, src2)] += np.dot(X[src2].T, X[src])
                
        return scipy.special.logsumexp(posteriors,axis=1).sum()
//  [python]
def hmm_forward(obs_loglik, trans_loglik):
  """Forward algorithm for a hidden Markov model."""
  n, k = obs_loglik.shape
  trans = np.exp(trans_loglik)

  loglik = 0.0

  l = obs_loglik[0] - np.log(k)
  c = np.log(np.sum(np.exp(l)))
  loglik += c

  for i in range(1, n):
    p = np.exp(l - c)
    l = np.log(np.dot(p, trans)) + obs_loglik[i]
    c = np.log(np.sum(np.exp(l)))
    loglik += c

  return loglik
//  [python]
def hmm_forward(obs_loglik, trans_loglik):
  """Forward algorithm for a hidden Markov model."""
  n, k = obs_loglik.shape
  trans = np.exp(trans_loglik)

  loglik = 0.0

  l = obs_loglik[0] - np.log(k)
  c = scipy.special.logsumexp(l)
  loglik += c

  for i in range(1, n):
    p = np.exp(l - c)
    l = np.log(np.dot(p, trans)) + obs_loglik[i]
    c = scipy.special.logsumexp(l)
    loglik += c

  return loglik
//  [python]
def logz_true(self):
        # true value of the integral
        n = self.ndim
        logmult = n / 2. * np.log(np.pi) - scipy.special.gammaln(n / 2. + 1)
        logvols = np.zeros(len(self.Rs))

        # volumes = pi^(n/2)/gamma(n/2+1) * (R_{k+1}^n - R_{k}^n)
        logvols[:-1] = logmult + n * np.log(
            self.Rs[1:]) + np.log1p(-(self.Rs[:-1] / self.Rs[1:])**n)
        # the last vol is (2s)^n - pi^(n/2)/gamma(n/2+1) * R[-1]^n
        logvols[-1] = n * np.log(
            2 * self.S) + np.log1p(-np.exp(logmult + n * np.log(self.Rs[-1] /
                                                                (2 * self.S))))
        logprior = -n * np.log(2 * self.S)
        logz = np.logaddexp.reduce(self.logAs + logvols) + logprior

        return logz
//  [python]
def logz_true(self):
        # true value of the integral
        n = self.ndim
        logmult = n / 2. * np.log(np.pi) - scipy.special.gammaln(n / 2. + 1)
        logvols = np.zeros(len(self.Rs))

        # volumes = pi^(n/2)/gamma(n/2+1) * (R_{k+1}^n - R_{k}^n)
        logvols[:-1] = logmult + n * np.log(
            self.Rs[1:]) + np.log1p(-(self.Rs[:-1] / self.Rs[1:])**n)
        # the last vol is (2s)^n - pi^(n/2)/gamma(n/2+1) * R[-1]^n
        logvols[-1] = n * np.log(
            2 * self.S) + np.log1p(-np.exp(logmult + n * np.log(self.Rs[-1] /
                                                                (2 * self.S))))
        logprior = -n * np.log(2 * self.S)
        logz = scipy.special.logsumexp(self.logAs + logvols) + logprior

        return logz
//  [python]
def makecontour(xytrue, indicator, inputs, net=None, like=None, istart=0, single=True):
  pp.figure(figsize=(12,8))

  if isinstance(like, types.FunctionType):
    like = like(inputs)
    
  for i in range(6):
    netinput = inputs[istart+i:(istart+i+1),:]
    
    with torch.no_grad():
      pars = net(netinput).detach().cpu().numpy().flatten()
      # xm, xe, xc = netmeanGn2(netinput, net=net)
      
    xs, ys = np.meshgrid(xmid, xmid, indexing='ij')    
    
    dx  = pars[0::6] - xs[:,:,np.newaxis]
    Fxx = pars[1::6]**2
  
    dy  = pars[2::6] - ys[:,:,np.newaxis]
    Fyy = pars[3::6]**2
  
    Fxy = np.arctan(pars[4::6]) / (0.5*math.pi) * pars[1::6] * pars[3::6]

    
    weight = torch.softmax(torch.from_numpy(pars[5::6]),dim=0).numpy()

    logmod = s.logsumexp(np.log(weight) - 0.5*(Fxx*dx*dx + Fyy*dy*dy + 2*Fxy*dx*dy) + 0.5*np.log(Fxx*Fyy - Fxy*Fxy), axis=2)
        
    q = np.exp(logmod)
    q /= np.sum(q)
    
    pmax = np.max(like[istart+i])
    qmax = np.max(q)
    vmax = min(pmax, qmax)
    
    pp.subplot(3,2,i+1)

    pp.plot([xytrue[istart+i,0]], [xytrue[istart+i,1]], 'ro')

    # note these are set-value posterior levels, not true set-mass-containing contours
    pp.contour(xs, ys, like[istart+i], colors='k', alpha=0.3, levels=[0.01*vmax,0.14*vmax,0.61*vmax])
    pp.contour(xs, ys, q,              colors='r', alpha=0.3, levels=[0.01*vmax,0.14*vmax,0.61*vmax])
//  [python]
def makecontour(xytrue, indicator, inputs, net=None, like=None, istart=0, single=True):
  pp.figure(figsize=(12,8))

  if isinstance(like, types.FunctionType):
    like = like(inputs)
    
  for i in range(6):
    netinput = inputs[istart+i:(istart+i+1),:]
    
    with torch.no_grad():
      pars = net(netinput).detach().cpu().numpy().flatten()
      # xm, xe, xc = netmeanGn2(netinput, net=net)
      
    xs, ys = np.meshgrid(xmid, xmid, indexing='ij')    
    
    dx  = pars[0::6] - xs[:,:,np.newaxis]
    Fxx = pars[1::6]**2
  
    dy  = pars[2::6] - ys[:,:,np.newaxis]
    Fyy = pars[3::6]**2
  
    Fxy = np.arctan(pars[4::6]) / (0.5*math.pi) * pars[1::6] * pars[3::6]

    
    weight = torch.softmax(torch.from_numpy(pars[5::6]),dim=0).numpy()

    logmod = scs.logsumexp(np.log(weight) - 0.5*(Fxx*dx*dx + Fyy*dy*dy + 2*Fxy*dx*dy) + 0.5*np.log(Fxx*Fyy - Fxy*Fxy), axis=2)
        
    q = np.exp(logmod)
    q /= np.sum(q)
    
    pmax = np.max(like[istart+i])
    qmax = np.max(q)
    vmax = min(pmax, qmax)
    
    pp.subplot(3,2,i+1)

    pp.plot([xytrue[istart+i,0]], [xytrue[istart+i,1]], 'ro')

    # note these are set-value posterior levels, not true set-mass-containing contours
    pp.contour(xs, ys, like[istart+i], colors='k', alpha=0.3, levels=[0.01*vmax,0.14*vmax,0.61*vmax])
    pp.contour(xs, ys, q,              colors='r', alpha=0.3, levels=[0.01*vmax,0.14*vmax,0.61*vmax])
//  [python]
def f_test(eta, word_count, beta_doc_kv, siginv=siginv_test, mu = mu_test):
    # 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_kv))
    ) - Ndoc * np.log(np.sum(np.exp(eta)))
    part2 = 0.5 * (eta[:-1] - mu).T @ siginv @ (eta[:-1] - mu)
    return np.float32(part2 - part1)
//  [python]
def f_test(eta, word_count, beta_doc_kv, siginv=siginv_test, mu = mu_test):
    # 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_kv))
    ) - Ndoc * scipy.special.logsumexp(eta)
    part2 = 0.5 * (eta[:-1] - mu).T @ siginv @ (eta[:-1] - mu)
    return np.float32(part2 - part1)
//  [python]
def log_prob(self, value):
        value = np.array(value)
        value = np.expand_dims(value, axis=-1)
        mlog_prob = Normal(self.locs, self.scales).log_prob(value)
        if self._validate_args:
            self._validate_sample(value)
        wlog_prob = np.log(self.weights) + mlog_prob
        return scipy.special.logsumexp(wlog_prob, axis=-1)
//  [python]
def compute_oscr(self, pred_k, pred_u, labels, nclasses):
        if pred_k.shape[1] == nclasses + 1:
            x1, x2 = pred_k[:, -1],  pred_u[:, -1]
        else:
            x1 = -1 * self.temperature * np.log(np.sum(np.exp(pred_k / self.temperature), axis=1))
            x2 = -1 * self.temperature * np.log(np.sum(np.exp(pred_u / self.temperature), axis=1))
            # x1, x2 = np.max(pred_k, axis=1), np.max(pred_u, axis=1)
        self.known_score = x1
        self.unknown_score = x2
        k_target = np.concatenate((np.ones(len(x1)), np.zeros(len(x2))), axis=0)
        u_target = np.concatenate((np.zeros(len(x1)), np.ones(len(x2))), axis=0)
        predict = np.concatenate((x1, x2), axis=0)
        n = len(predict)

        # Cutoffs are of prediction values
        
        CCR = [0 for x in range(n+2)]
        FPR = [0 for x in range(n+2)] 

        auxa = np.random.random(predict.size)
        idx = np.lexsort((auxa, predict))
        # idx = predict.argsort()

        s_k_target = k_target[idx]
        s_u_target = u_target[idx]
        for k in range(n-1):
            CC = s_u_target[k+1:].sum()
            FP = s_k_target[k:].sum()

            # True	Positive Rate
            CCR[k] = float(CC) / float(len(x2))
            # False Positive Rate
            FPR[k] = float(FP) / float(len(x1))

        CCR[n] = 0.0
        FPR[n] = 0.0
        CCR[n+1] = 1.0
        FPR[n+1] = 1.0

        # Positions of ROC curve (FPR, TPR)
        ROC = sorted(zip(FPR, CCR), reverse=True)

        OSCR = 0

        # Compute AUROC Using Trapezoidal Rule
        for j in range(n+1):
            h =   ROC[j][0] - ROC[j+1][0]
            w =  (ROC[j][1] + ROC[j+1][1]) / 2.0

            OSCR = OSCR + h*w

        return OSCR, ROC
//  [python]
def compute_oscr(self, pred_k, pred_u, labels, nclasses):
        if pred_k.shape[1] == nclasses + 1:
            x1, x2 = pred_k[:, -1],  pred_u[:, -1]
        else:
            x1 = -1 * self.temperature * logsumexp(pred_k / self.temperature, axis=1)
            x2 = -1 * self.temperature * logsumexp(pred_u / self.temperature, axis=1)
            # x1, x2 = np.max(pred_k, axis=1), np.max(pred_u, axis=1)
        self.known_score = x1
        self.unknown_score = x2
        k_target = np.concatenate((np.ones(len(x1)), np.zeros(len(x2))), axis=0)
        u_target = np.concatenate((np.zeros(len(x1)), np.ones(len(x2))), axis=0)
        predict = np.concatenate((x1, x2), axis=0)
        n = len(predict)

        # Cutoffs are of prediction values
        
        CCR = [0 for x in range(n+2)]
        FPR = [0 for x in range(n+2)] 

        auxa = np.random.random(predict.size)
        idx = np.lexsort((auxa, predict))
        # idx = predict.argsort()

        s_k_target = k_target[idx]
        s_u_target = u_target[idx]
        for k in range(n-1):
            CC = s_u_target[k+1:].sum()
            FP = s_k_target[k:].sum()

            # True	Positive Rate
            CCR[k] = float(CC) / float(len(x2))
            # False Positive Rate
            FPR[k] = float(FP) / float(len(x1))

        CCR[n] = 0.0
        FPR[n] = 0.0
        CCR[n+1] = 1.0
        FPR[n+1] = 1.0

        # Positions of ROC curve (FPR, TPR)
        ROC = sorted(zip(FPR, CCR), reverse=True)

        OSCR = 0

        # Compute AUROC Using Trapezoidal Rule
        for j in range(n+1):
            h =   ROC[j][0] - ROC[j+1][0]
            w =  (ROC[j][1] + ROC[j+1][1]) / 2.0

            OSCR = OSCR + h*w

        return OSCR, ROC
//  [python]
def _compute_kappa(
    ln_z: numpy.ndarray, ln_q: numpy.ndarray, factor: numpy.ndarray, n: int
) -> tuple[numpy.ndarray, numpy.ndarray]:
    ln_z = numpy.insert(ln_z, 0, 0.0)

    ln_q_ij_sum = scipy.special.logsumexp(a=ln_q - ln_z, b=factor, axis=1)
    kappa = ln_q_ij_sum.sum() / n + (factor * ln_z).sum()

    w = factor * _compute_weights(ln_z, ln_q, factor)
    grad = -w[:, 1:].sum(axis=0) / n + factor[1:]

    return kappa, grad
//  [python]
def test_logsumexp(axis: Optional[int], keepdims: bool):
    rng = np.random.default_rng(0)

    n_sources, n_channels = 4, 3
    n_frames = 8
    shape = (n_sources, n_frames, n_channels, n_channels)

    X = rng.random(shape)

    Y = logsumexp(X, axis=axis, keepdims=keepdims)
    Y_scipy = scipy.special.logsumexp(X, axis=axis, keepdims=keepdims)

    assert np.allclose(Y, Y_scipy)
//  [python]
def get_score(self, defn):
        empty_prompt = self.list_to_prompt([], False)
        if len(empty_prompt) > 0 and empty_prompt[-1] == " ":
            empty_prompt = empty_prompt[:-1]
        empty_input_ids = self.tokenizer(empty_prompt, return_tensors="pt")

        tokeninfos = []
        prompt = self.list_to_prompt(defn, True)
        input_ids = self.tokenizer(prompt, return_tensors="pt").input_ids.to(args.device)

        if len(input_ids[0]) > self.model.config.max_position_embeddings:
            raise Exception("prompt has too many tokens: {}".format(len(input_ids[0])))

        entropy = 0
        output = self.model(input_ids)
        logits = output.logits[0].to("cpu");
        for ii in range(len(empty_input_ids[0]), len(input_ids[0])):
            total = 0
            tokenid = input_ids[0][ii].item()
            tokeninfo = {}
            tokeninfo["token"] = self.tokenizer.decode([tokenid])

            tok_logits = logits[ii-1]

            # use scipy.special.logsumexp to avoid possible overflow
            # (commonly happens with smaller models like pythia-70M)
            prob = np.exp(tok_logits[tokenid] - np.logaddexp.reduce(tok_logits))

            new_entropy = math.log2(1.0 / prob)
            entropy += new_entropy
            tokeninfo["score"] = new_entropy
            tokeninfos.append(tokeninfo)

        return entropy, tokeninfos
//  [python]
def get_score(self, defn):
        empty_prompt = self.list_to_prompt([], False)
        if len(empty_prompt) > 0 and empty_prompt[-1] == " ":
            empty_prompt = empty_prompt[:-1]
        empty_input_ids = self.tokenizer(empty_prompt, return_tensors="pt")

        tokeninfos = []
        prompt = self.list_to_prompt(defn, True)
        input_ids = self.tokenizer(prompt, return_tensors="pt").input_ids.to(args.device)

        if len(input_ids[0]) > self.model.config.max_position_embeddings:
            raise Exception("prompt has too many tokens: {}".format(len(input_ids[0])))

        entropy = 0
        output = self.model(input_ids)
        logits = output.logits[0].to("cpu");
        for ii in range(len(empty_input_ids[0]), len(input_ids[0])):
            total = 0
            tokenid = input_ids[0][ii].item()
            tokeninfo = {}
            tokeninfo["token"] = self.tokenizer.decode([tokenid])

            tok_logits = logits[ii-1]

            # use scipy.special.logsumexp to avoid possible overflow
            # (commonly happens with smaller models like pythia-70M)
            prob = np.exp(tok_logits[tokenid] - scipy.special.logsumexp(tok_logits))

            new_entropy = math.log2(1.0 / prob)
            entropy += new_entropy
            tokeninfo["score"] = new_entropy
            tokeninfos.append(tokeninfo)

        return entropy, tokeninfos
//  [python]
def optimize_eta(eta, word_count, beta_doc_kv, mu=mu_test, siginv=siginv_test):
    def f_test(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))
        part1 = np.dot(
            word_count, (eta.max() + np.log(np.exp(eta - eta.max()) @ beta_doc_kv))
        ) - Ndoc * np.log(np.sum(np.exp(eta)))
        part2 = 0.5 * (eta[:-1] - mu).T @ siginv @ (eta[:-1] - mu)
        return np.float32(part2 - part1)

    def df_test(eta, word_count, beta_doc_kv):
        """gradient for the objective of the variational update q(etas)"""
        # precomputation
        eta = np.insert(eta, K - 1, 0)
        # 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)
        part1 = np.delete(
            beta_doc_kv @ (word_count / np.sum(beta_doc_kv.T, axis=1))
            - np.sum(word_count) / np.sum(np.exp(eta)),
            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_test, x0=eta, args=(word_count, beta_doc_kv), jac=df_test, method="BFGS", options={'maxiter': 500, 'gtol': 1e-01, 'eps':1},
    ).x
//  [python]
def optimize_eta(eta, word_count, beta_doc_kv, mu=mu_test, siginv=siginv_test):
    def f_test(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))
        part1 = np.dot(
            word_count, (eta.max() + np.log(np.exp(eta - eta.max()) @ beta_doc_kv))
        ) - Ndoc * scipy.special.logsumexp(eta)
        part2 = 0.5 * (eta[:-1] - mu).T @ siginv @ (eta[:-1] - mu)
        return np.float32(part2 - part1)

    def df_test(eta, word_count, beta_doc_kv):
        """gradient for the objective of the variational update q(etas)"""
        # precomputation
        eta = np.insert(eta, K - 1, 0)
        # 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)
        part1 = np.delete(
            beta_doc_kv @ (word_count / np.sum(beta_doc_kv.T, axis=1))
            - np.sum(word_count) / np.sum(np.exp(eta)),
            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_test, x0=eta, args=(word_count, beta_doc_kv), jac=df_test, method="BFGS", options={'maxiter': 500, 'gtol': 1e-01, 'eps':1},
    ).x
//  [python]
def test_logsumexp(a, b, dim, keepdim):
    from scipy.special import logsumexp as s

    actual = smee.utils.logsumexp(a, dim, keepdim, b)
    expected = torch.tensor(
        s(a.numpy(), dim, b if b is None else b.numpy(), keepdim)
    )

    assert actual.shape == expected.shape
    assert torch.allclose(actual.double(), expected.double())
//  [python]
def test_logsumexp(a, b, dim, keepdim):
    from scipy.special import logsumexp

    actual = smee.utils.logsumexp(a, dim, keepdim, b)
    expected = torch.tensor(
        logsumexp(a.numpy(), dim, b if b is None else b.numpy(), keepdim)
    )

    assert actual.shape == expected.shape
    assert torch.allclose(actual.double(), expected.double())
//  [python]
def test_logsumexp_with_sign():
    import scipy.special as s

    a = torch.tensor([1.0, 2.0, 3.0])
    b = torch.tensor(-2.0)

    actual, actual_sign = smee.utils.logsumexp(a, -1, True, b, return_sign=True)
    expected, expected_sign = torch.tensor(
        s.logsumexp(a.numpy(), -1, b.numpy(), True, return_sign=True)
    )

    assert actual.shape == expected.shape
    assert torch.allclose(actual.double(), expected.double())

    assert actual_sign.shape == expected_sign.shape
    assert torch.allclose(actual_sign.double(), expected_sign.double())
//  [python]
def test_logsumexp_with_sign():
    from scipy.special import logsumexp

    a = torch.tensor([1.0, 2.0, 3.0])
    b = torch.tensor(-2.0)

    actual, actual_sign = smee.utils.logsumexp(a, -1, True, b, return_sign=True)
    expected, expected_sign = torch.tensor(
        logsumexp(a.numpy(), -1, b.numpy(), True, return_sign=True)
    )

    assert actual.shape == expected.shape
    assert torch.allclose(actual.double(), expected.double())

    assert actual_sign.shape == expected_sign.shape
    assert torch.allclose(actual_sign.double(), expected_sign.double())
//  [python]
def _log_density(self, stimulus):
        log_densities = []
        for i, model in enumerate(self.models):
            log_density = model.log_density(stimulus).copy()
            log_density += np.log(self.weights[i])
            log_densities.append(log_density)

        log_density = np.logaddexp.reduce(log_densities, axis=0)

        if self.check_norm:
            np.testing.assert_allclose(np.exp(log_density).sum(), 1.0, rtol=1e-7)

        if not log_density.shape == (stimulus.shape[0], stimulus.shape[1]):
            raise ValueError('wrong density shape in mixture model! stimulus shape: ({}, {}), density shape: {}'.format(stimulus.shape[0], stimulus.shape[1], log_density.shape))
        return log_density
//  [python]
def _log_density(self, stimulus):
        log_densities = []
        for i, model in enumerate(self.models):
            log_density = model.log_density(stimulus).copy()
            log_density += np.log(self.weights[i])
            log_densities.append(log_density)

        log_density = logsumexp(log_densities, axis=0)

        if self.check_norm:
            np.testing.assert_allclose(np.exp(log_density).sum(), 1.0, rtol=1e-7)

        if not log_density.shape == (stimulus.shape[0], stimulus.shape[1]):
            raise ValueError('wrong density shape in mixture model! stimulus shape: ({}, {}), density shape: {}'.format(stimulus.shape[0], stimulus.shape[1], log_density.shape))
        return log_density
//  [python]
def conditional_log_density(self, stimulus, x_hist, y_hist, t_hist, attributes=None, out=None):
        log_densities = []
        for i, model in enumerate(self.models):
            log_density = model.conditional_log_density(stimulus, x_hist, y_hist, t_hist, attributes=attributes).copy()
            log_density += np.log(self.weights[i])
            log_densities.append(log_density)

        log_density = np.max(log_densities, axis=0) + np.log(np.sum(np.exp(log_densities - np.max(log_densities, axis=0)), axis=0))
        if self.check_norm:
            np.testing.assert_allclose(np.exp(log_density).sum(), 1.0, rtol=1e-7)
        if not log_density.shape == (stimulus.shape[0], stimulus.shape[1]):
            raise ValueError('wrong density shape in mixture model! stimulus shape: ({}, {}), density shape: {}'.format(stimulus.shape[0], stimulus.shape[1], log_density.shape))
        return log_density
//  [python]
def conditional_log_density(self, stimulus, x_hist, y_hist, t_hist, attributes=None, out=None):
        log_densities = []
        for i, model in enumerate(self.models):
            log_density = model.conditional_log_density(stimulus, x_hist, y_hist, t_hist, attributes=attributes).copy()
            log_density += np.log(self.weights[i])
            log_densities.append(log_density)

        log_density = logsumexp(log_densities, axis=0)
        if self.check_norm:
            np.testing.assert_allclose(np.exp(log_density).sum(), 1.0, rtol=1e-7)
        if not log_density.shape == (stimulus.shape[0], stimulus.shape[1]):
            raise ValueError('wrong density shape in mixture model! stimulus shape: ({}, {}), density shape: {}'.format(stimulus.shape[0], stimulus.shape[1], log_density.shape))
        return log_density
//  [python]
def _log_density(self, stimulus):
        smap = self.parent_model.log_density(stimulus)

        target_shape = (stimulus.shape[0],
                        stimulus.shape[1])

        if smap.shape != target_shape:
            if self.verbose:
                print("Resizing saliency map", smap.shape, target_shape)
            x_factor = target_shape[1] / smap.shape[1]
            y_factor = target_shape[0] / smap.shape[0]

            smap = zoom(smap, [y_factor, x_factor], order=1, mode='nearest')

            smap -= np.log(np.sum(np.exp(smap)))

            assert smap.shape == target_shape

        return smap
//  [python]
def _log_density(self, stimulus):
        smap = self.parent_model.log_density(stimulus)

        target_shape = (stimulus.shape[0],
                        stimulus.shape[1])

        if smap.shape != target_shape:
            if self.verbose:
                print("Resizing saliency map", smap.shape, target_shape)
            x_factor = target_shape[1] / smap.shape[1]
            y_factor = target_shape[0] / smap.shape[0]

            smap = zoom(smap, [y_factor, x_factor], order=1, mode='nearest')

            smap -= logsumexp(smap)

            assert smap.shape == target_shape

        return smap
//  [python]
def conditional_log_density(self, stimulus, x_hist, y_hist, t_hist, attributes=None, out=None):
        smap = self.parent_model.conditional_log_density(stimulus, x_hist, y_hist, t_hist, attributes=attributes, out=out)

        target_shape = (stimulus.shape[0],
                        stimulus.shape[1])

        if smap.shape != target_shape:
            if self.verbose:
                print("Resizing saliency map", smap.shape, target_shape)
            x_factor = target_shape[1] / smap.shape[1]
            y_factor = target_shape[0] / smap.shape[0]

            smap = zoom(smap, [y_factor, x_factor], order=1, mode='nearest')

            smap -= np.max(smap) #deprecated API
            smap -= np.log(np.sum(np.exp(smap - np.max(smap)))) #deprecated API

            assert smap.shape == target_shape

        return smap
//  [python]
def conditional_log_density(self, stimulus, x_hist, y_hist, t_hist, attributes=None, out=None):
        smap = self.parent_model.conditional_log_density(stimulus, x_hist, y_hist, t_hist, attributes=attributes, out=out)

        target_shape = (stimulus.shape[0],
                        stimulus.shape[1])

        if smap.shape != target_shape:
            if self.verbose:
                print("Resizing saliency map", smap.shape, target_shape)
            x_factor = target_shape[1] / smap.shape[1]
            y_factor = target_shape[0] / smap.shape[0]

            smap = zoom(smap, [y_factor, x_factor], order=1, mode='nearest')

            smap -= logsumexp(smap)

            assert smap.shape == target_shape

        return smap
//  [python]
def cond_log_prob(
        self,
        inputs: Union[str, List[str]],
        targets: Union[List[str], List[List[str]]],
        absolute_normalization: Optional[bool] = False,
    ) -> Union[List[float], List[List[float]]]:
        """Computes conditional log probabilities of targets given inputs.

        Args:
          inputs: input string or list of strings
          targets: target strings for each input
          absolute_normalization: When True, the function returns the log
            probability of unconstrained generation or the target sequence. When
            False (default), log probabilities are normalized so that the
            probabilities of generating `targets` sum to 1. Note that setting
            `absolute_normalization` to True restricts the class of models that can
            be evaluated to those that can assign absolute probabilities to
            sequences.

        Returns:
          conditional log probabilities for inputs/targets
        """
        absolute_scores = self.model.cond_log_prob(
            inputs=inputs,
            targets=targets,
            absolute_normalization=True,
        )

        if isinstance(inputs, list):
            sample_inputs = inputs
            sample_targets = targets
            absolute_scores = absolute_scores
        else:
            sample_inputs = [inputs]
            sample_targets = [targets]
            absolute_scores = [absolute_scores]

        normalized_scores = [
            list(score_row - s.logsumexp(score_row))
            for score_row in absolute_scores
        ]

        if absolute_normalization:
            sample_scores = absolute_scores
        else:
            sample_scores = normalized_scores

        samples = []

        for input_str, target_list, score_list, abs_scores, norm_scores in zip(
            sample_inputs,
            sample_targets,
            sample_scores,
            absolute_scores,
            normalized_scores,
        ):
            # if this is a m/c sample, we will update later for json tasks
            samples.append(
                results_api.ScoringSample(
                    input=input_str,
                    targets=target_list,
                    scores=score_list,
                    absolute_scores=abs_scores,
                    normalized_scores=norm_scores,
                )
            )

        # we create these as ScoringQueries, but if this is a m/c
        # query, we will change this later for json tasks
        self.queries.append(
            results_api.ScoringQuery(
                function="cond_log_prob",
                absolute_normalization=absolute_normalization,
                samples=samples,
            )
        )

        if isinstance(inputs, list):
            return sample_scores
        else:
            return sample_scores[0]
//  [python]
def cond_log_prob(
        self,
        inputs: Union[str, List[str]],
        targets: Union[List[str], List[List[str]]],
        absolute_normalization: Optional[bool] = False,
    ) -> Union[List[float], List[List[float]]]:
        """Computes conditional log probabilities of targets given inputs.

        Args:
          inputs: input string or list of strings
          targets: target strings for each input
          absolute_normalization: When True, the function returns the log
            probability of unconstrained generation or the target sequence. When
            False (default), log probabilities are normalized so that the
            probabilities of generating `targets` sum to 1. Note that setting
            `absolute_normalization` to True restricts the class of models that can
            be evaluated to those that can assign absolute probabilities to
            sequences.

        Returns:
          conditional log probabilities for inputs/targets
        """
        absolute_scores = self.model.cond_log_prob(
            inputs=inputs,
            targets=targets,
            absolute_normalization=True,
        )

        if isinstance(inputs, list):
            sample_inputs = inputs
            sample_targets = targets
            absolute_scores = absolute_scores
        else:
            sample_inputs = [inputs]
            sample_targets = [targets]
            absolute_scores = [absolute_scores]

        normalized_scores = [
            list(score_row - scipy.special.logsumexp(score_row))
            for score_row in absolute_scores
        ]

        if absolute_normalization:
            sample_scores = absolute_scores
        else:
            sample_scores = normalized_scores

        samples = []

        for input_str, target_list, score_list, abs_scores, norm_scores in zip(
            sample_inputs,
            sample_targets,
            sample_scores,
            absolute_scores,
            normalized_scores,
        ):
            # if this is a m/c sample, we will update later for json tasks
            samples.append(
                results_api.ScoringSample(
                    input=input_str,
                    targets=target_list,
                    scores=score_list,
                    absolute_scores=abs_scores,
                    normalized_scores=norm_scores,
                )
            )

        # we create these as ScoringQueries, but if this is a m/c
        # query, we will change this later for json tasks
        self.queries.append(
            results_api.ScoringQuery(
                function="cond_log_prob",
                absolute_normalization=absolute_normalization,
                samples=samples,
            )
        )

        if isinstance(inputs, list):
            return sample_scores
        else:
            return sample_scores[0]