class

Multinomial

extendsDistribution

Multinomial(total_count: Tensor | int = 1, probs: Tensor | None = None, logits: Tensor | None = None, validate_args: bool | None = None)

source

Multinomial distribution over $K$ categories.

Multivariate generalisation of the lucid.distributions.Bernoulli /Binomial: the joint distribution of the counts $(k_1, \ldots, k_K)$ obtained when $n = \text{total\_count}$ independent draws are made from a Categorical with probabilities $(p_1, \ldots, p_K)$ . The $k_i$ are non-negative integers summing to $n$ .

The event dimension is the last axis of probs / logits; all preceding axes form the batch shape.

Parameters

total_countint or Tensor= 1

Total number of trials

n \geq 0

. Default 1 (a one-hot Categorical sample).

probsTensor= None

Unnormalised category probabilities of shape (..., K) — normalised internally to sum to 1 along the last axis. Mutually exclusive with logits.

logitsTensor= None

Unnormalised log-probabilities of shape (..., K), converted via softmax. Mutually exclusive with probs.

validate_argsbool= None

If True, validate parameter constraints at construction time.

Notes

Probability mass function (joint over the count vector $\mathbf{k}$ with $\sum_i k_i = n$ ):

P(\mathbf{K} = \mathbf{k} \mid n, \mathbf{p}) = \binom{n}{k_1, \ldots, k_K} \prod_{i=1}^{K} p_i^{k_i}, \qquad \binom{n}{k_1, \ldots, k_K} = \frac{n!}{\prod_i k_i!}

Moments (per category):

\mathbb{E}[K_i] = n p_i, \qquad \mathrm{Var}[K_i] = n p_i (1 - p_i), \qquad \mathrm{Cov}[K_i, K_j] = -n p_i p_j \;\; (i \neq j)

Special cases:

$K = 2$ → $\mathrm{Binomial}(n, p_1)$
$n = 1$ → one-hot lucid.distributions.Categorical

Conjugate prior: lucid.distributions.Dirichlet — observing counts $\mathbf{k}$ updates Dirichlet(α) → Dirichlet(α + k).

Sampling is non-reparameterised (has_rsample = False) and implemented by summing $n$ one-hot Categorical draws.

Examples

>>> import lucid
>>> from lucid.distributions import Multinomial
>>> d = Multinomial(total_count=10, probs=lucid.tensor([0.2, 0.3, 0.5]))
>>> d.mean  # n * p
Tensor([2., 3., 5.])
>>> d.sample((4,))
Tensor([...])

Methods (7)

dunder

init

→None

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

source

Initialise a Multinomial distribution.

Parameters

total_countTensor | int= 1

Total number of draws

n \geq 0

. Must be a scalar. Default is 1 (reduces to a one-hot Categorical).

probsTensor | None= None

Unnormalised probability vector(s) of shape (..., K). Normalised to sum to 1 internally. Mutually exclusive with logits.

logitsTensor | None= None

Unnormalised log-probability vector(s) of shape (..., K). Converted to probabilities via softmax. Mutually exclusive with probs.

validate_argsbool | None= None

If True, validate parameter constraints at construction time.

prop

support

→Constraint

support: Constraint

source

Support of the distribution: non-negative integers.

Returns

Constraint

nonnegative_integer — the multinomial event count vector lives in $\mathbb{Z}_{\geq 0}^K$ , subject to the additional constraint that the counts sum to total_count.

prop

total_count

→Tensor

total_count: Tensor

source

Total number of trials $n$ per Multinomial draw.

Returns

Tensor

The integer-valued n parameter broadcast over the batch shape. Each independent Multinomial sums to this many trials across the K categories.

prop

mean

→Tensor

mean: Tensor

source

n · p (element-wise).

prop

variance

→Tensor

variance: Tensor

source

n · p · (1 − p) (element-wise).

log_prob

→Tensor

log_prob(value: Tensor)

source

log C(n; k₁,…,kK) + Σ kᵢ log pᵢ.

sample

→Tensor

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

source

Draw samples by summing one-hot Categorical draws.