class

NegativeBinomial

extendsDistribution

NegativeBinomial(total_count: Tensor | float, probs: Tensor | float | None = None, logits: Tensor | float | None = None, validate_args: bool | None = None)

source

Negative Binomial distribution over the number of failures before r successes.

NegativeBinomial(total_count=r, probs=p) models the count of failures $k$ before achieving $r$ successes in a sequence of independent Bernoulli trials each with success probability $1 - p$ .

Equivalently, this is a Gamma-Poisson compound: draw $\lambda \sim \operatorname{Gamma}(r, (1-p)/p)$ , then $X \sim \operatorname{Poisson}(\lambda)$ . This interpretation extends the distribution to real-valued total_count (the generalised Negative Binomial) and forms the basis of the sampler.

Parameters

total_countTensor | float

Dispersion / number of successes

r > 0

. May be non-integer for the generalised version.

probsTensor | float | None= None

Probability of failure

p \in (0, 1)

. Mutually exclusive with logits.

logitsTensor | float | None= None

Log-odds of failure

l = \log(p / (1-p)) \in \mathbb{R}

. Mutually exclusive with probs.

validate_argsbool | None= None

If True, validate parameter constraints at construction time.

Attributes

total_countTensor

Dispersion parameter

r

probsTensor

Failure probability (present when constructed with probs).

logitsTensor

Log-odds of failure (present when constructed with logits).

Notes

PMF:

P(X = k) = \binom{k + r - 1}{k} (1-p)^r p^k, \quad k \in \{0, 1, 2, \ldots\}

Log-PMF (via lgamma for numerical stability):

\log P(X = k) = \log\Gamma(k+r) - \log\Gamma(r) - \log\Gamma(k+1) + r \log(1-p) + k \log p

Moments:

Mean: $E[X] = r p / (1-p)$
Variance: $\operatorname{Var}[X] = r p / (1-p)^2$

Sampler: uses the Gamma-Poisson representation. A standard-Gamma variate with concentration $r$ is drawn via rejection sampling, scaled by $p/(1-p)$ , and passed to the Poisson sampler. This produces exact (non-approximate) samples for all real $r > 0$ .

Examples

>>> import lucid
>>> from lucid.distributions import NegativeBinomial
>>> dist = NegativeBinomial(total_count=5.0, probs=0.4)
>>> samples = dist.sample((100,))
>>> samples.shape
(100,)
>>> dist.mean

Methods (5)

dunder

init

→None

__init__(total_count: Tensor | float, probs: Tensor | float | None = None, logits: Tensor | float | None = None, validate_args: bool | None = None)

source

Initialise a Negative Binomial distribution.

Parameters

total_countTensor | float

Dispersion / number of successes

r > 0

. May be non-integer for the generalised version.

probsTensor | float | None= None

Probability of failure

p \in (0, 1)

. Mutually exclusive with logits.

logitsTensor | float | None= None

Log-odds of failure

l = \log(p / (1-p))

. Mutually exclusive with probs.

validate_argsbool | None= None

If True, validate parameter constraints at construction time.

Raises

ValueError

If both or neither of probs and logits are provided.

prop

mean

→Tensor

mean: Tensor

source

Mean of the Negative Binomial: $E[X] = r p / (1-p)$ .

Returns

Tensor

Mean values of shape batch_shape.

prop

variance

→Tensor

variance: Tensor

source

Variance of the Negative Binomial: $\operatorname{Var}[X] = r p / (1-p)^2$ .

Returns

Tensor

Variance values of shape batch_shape.

sample

→Tensor

sample(sample_shape: tuple[int, ...] = ())

source

Draw samples via the Gamma-Poisson compound representation.

Draws $\lambda \sim \operatorname{Gamma}(r, (1-p)/p)$ then $X \sim \operatorname{Poisson}(\lambda)$ . This produces exact (non-approximate) samples for all real $r > 0$ .

Parameters

sample_shapetuple[int, ...]= ()

Leading shape of the output sample batch.

Returns

Tensor

Non-negative integer samples of shape (*sample_shape, *batch_shape).

log_prob

→Tensor

log_prob(value: Tensor)

source

Log-probability of counts under the Negative Binomial distribution.

\log P(X = k) = \log\Gamma(k+r) - \log\Gamma(r) - \log\Gamma(k+1) + r \log(1-p) + k \log p

Parameters

valueTensor

Non-negative integer counts

k \geq 0

Returns

Tensor

Log-probability values of the same shape as value.