def logpmf(self, x, n, p):
x, n, p = _promote_dtypes(x, n, p)
if self.is_logits:
return gammaln(n + 1) + np.sum(x * p - gammaln(x + 1),
axis=-1) - n * logsumexp(p, axis=-1)
else:
return gammaln(n + 1) + np.sum(xlogy(x, p) - gammaln(x + 1),
axis=-1)
def _logpmf(self, x, n, p):
x, n, p = _promote_dtypes(x, n, p)
combiln = gammaln(n + 1) - (gammaln(x + 1) + gammaln(n - x + 1))
if self.is_logits:
# TODO: move this implementation to PyTorch if it does not get non-continuous problem
# In PyTorch, k * logit - n * log1p(e^logit) get overflow when logit is a large
# positive number. In that case, we can reformulate into
# k * logit - n * log1p(e^logit) = k * logit - n * (log1p(e^-logit) + logit)
# = k * logit - n * logit - n * log1p(e^-logit)
# More context: https://github.com/pytorch/pytorch/pull/15962/
return combiln + x * p - (n * jnp.clip(p, 0) + xlog1py(n, jnp.exp(-jnp.abs(p))))
else:
return combiln + xlogy(x, p) + xlog1py(n - x, -p)
def _ppf(self, q, a, b):
q, a, b = _promote_dtypes(q, a, b)
# XXX: consider to use norm._ppf(q * norm._cdf(b) + norm._cdf(a) * (1.0 - q))
# when a, b < 0
ppf = norm._isf(q * norm._sf(b) + norm._sf(a) * (1.0 - q))
return ppf
def _logpdf(self, x, a, b):
x, a, b = _promote_dtypes(x, a, b)
# XXX: consider to use norm._cdf(b) - norm._cdf(a) when a, b < 0
delta = norm._sf(a) - norm._sf(b)
return norm._logpdf(x) - np.log(delta)
def _logpmf(self, x, mu):
x, mu = _promote_dtypes(x, mu)
Pk = xlogy(x, mu) - gammaln(x + 1) - mu
return Pk
def _logpmf(self, x, n, p):
k = np.floor(x)
n, p = _promote_dtypes(n, p)
combiln = (gammaln(n + 1) - (gammaln(k + 1) + gammaln(n - k + 1)))
return combiln + xlogy(k, p) + xlog1py(n - k, -p)
def _logpmf(self, x, n, p):
n, p, x = _promote_dtypes(n, p, x)
return gammaln(n + 1) + np.sum(xlogy(x, p) - gammaln(x + 1), axis=-1)
