distributions

Bernoulli distribution over $\{0, 1\}$.

Geometric distribution: number of failures before the first success.

Categorical distribution — a discrete distribution over K labelled outcomes.

Categorical distribution with one-hot encoded samples.

Fisher-Snedecor F-distribution — ratio of two scaled Chi-squared variates.

Half-Cauchy distribution — the absolute value of a zero-location Cauchy.

Half-Normal distribution — the absolute value of a zero-mean Normal.

Pareto distribution — a heavy-tailed power-law family on $[x_m, \infty)$.

Weibull distribution — a flexible family for lifetime and survival analysis.

Binomial distribution over the number of successes in n independent trials.

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

Poisson distribution over the non-negative integers.

Abstract base for a probability distribution.

Mix-in for exponential-family distributions.

Cauchy distribution on $\mathbb{R}$ — heavy-tailed, all moments undefined.

Exponential distribution on $[0, \infty)$.

Laplace (double-exponential) distribution on $\mathbb{R}$.

Beta distribution on the open interval $(0, 1)$.

Chi-squared distribution with $k$ degrees of freedom.

Dirichlet distribution on the $K$-simplex.

Gamma distribution on $(0, \infty)$ — shape/rate parameterisation.

Re-interpret rightmost batch dimensions of a base distribution as event dims.

Finite mixture model where all components share the same distribution family.

Multivariate Normal (Gaussian) distribution in $\mathbb{R}^D$.

Log-Normal distribution: $X = \exp(Y)$ where $Y \sim \mathcal{N}(\mu, \sigma^2)$.

Univariate Gaussian (Normal) distribution $\mathcal{N}(\mu, \sigma^2)$.

Concrete (Gumbel-sigmoid) relaxation of the Bernoulli distribution.

Concrete distribution — a continuous relaxation of OneHotCategorical.

Student's t-distribution with location, scale, and degrees of freedom.

Element-wise absolute value $y = |x|$ (folded, not bijective).

Element-wise affine bijection $y = \mathrm{loc} + \mathrm{scale}\cdot x$.

Apply different transforms to contiguous partitions along an axis.

Function composition of a list of transforms — left-to-right.

Bijection from $\mathbb{R}^{d(d-1)/2}$ to Cholesky factors of correlation matrices.

Probability-integral transform $y = F(x)$ via a base distribution's CDF.

Element-wise exponential bijection $y = e^x$.

Reinterpret n trailing batch dimensions of an inner transform as event dims.

Bijection mapping an unconstrained matrix to a positive-diagonal Cholesky factor.

Element-wise power bijection $y = x^{\mathrm{exponent}}$.

Pure-shape reinterpretation of the event shape $y = \mathrm{reshape}(x)$.

Element-wise sigmoid bijection $y = \sigma(x) = 1/(1 + e^{-x})$.

Softmax transform mapping $\mathbb{R}^K \to \Delta^{K-1}$.

Apply a list of transforms to indexed slices along a stack axis.

Logistic stick-breaking bijection $\mathbb{R}^{K-1} \to \Delta^{K-1}$.

Element-wise hyperbolic tangent bijection $y = \tanh(x)$.

Abstract bijection between two measurable spaces.

Pushforward of a base distribution through a (composite) bijector.

Continuous Bernoulli distribution on $[0, 1]$.

Gumbel (Type-I extreme value) distribution on $\mathbb{R}$.

Inverse-Gamma distribution on $(0, \infty)$.

Kumaraswamy distribution on $(0, 1)$.

Multinomial distribution over $K$ categories.

LKJ distribution over Cholesky factors of correlation matrices.

Wishart distribution over symmetric positive-definite matrices.

Continuous Uniform distribution on the half-open interval $[a, b)$.