Module Norn.Dist

The type for distributions over values 'x whose log densities have element type 'f. 'x is (float, 'f, 'd) Nx.t for continuous families, 'd Nx.int32_t for counts, 'd Nx.int64_t for categories and 'd Nx.bool_t for bernoulli.

The type for the values of a distribution: a witness of 'x.

normal ~loc ~scale has density exp (-z²/2) / (scale sqrt (2π)), z = (x - loc) / scale. scale is in (0, inf) and loc finite.

half_normal ~scale is the distribution of |x|, x normal with location 0, on (0, inf).

lognormal ~loc ~scale is the distribution of exp x, x normal.

student_t ~df ~loc ~scale is Student's t with df degrees of freedom, in (0, inf), located and scaled.

cauchy ~loc ~scale has density 1 / (π scale (1 + z²)).

half_cauchy ~scale is the distribution of |x|, x Cauchy with location 0, on (0, inf).

laplace ~loc ~scale has density exp (-|z|) / (2 scale).

logistic ~loc ~scale has density exp (-z) / (scale (1 + exp (-z))²).

exponential ~rate has density rate exp (-rate x) on (0, inf).

gamma ~concentration ~rate has density proportional to x^(α - 1) exp (-rate x) on (0, inf), α the concentration.

inverse_gamma ~concentration ~scale is the distribution of 1 / x, x gamma with rate scale.

beta ~a ~b has density proportional to x^(a - 1) (1 - x)^(b - 1) on (0, 1).

uniform ~low ~high has density 1 / (high - low) on (low, high). Each element of high - low is in (0, inf).

dirichlet ~concentration is over vectors along the last axis of concentration, of positive components summing to one, with density proportional to prod_k x_k^(α_k - 1). A matrix of concentrations is independent rows. A vector whose sum is within 4 K ε of one, K its length, is read as x / Σ x, the point of the simplex on its ray, so the rounding that keeps it off the simplex does not move its density. Draws are Nx.Rng.dirichlet's, made at float32 for a narrower dtype, each component at least the dtype's least positive value.

Raises Invalid_argument if the last axis has fewer than two components.

mvn ~loc ~scale_tril is the multivariate normal of vectors along the last axis, with mean loc and covariance L Lᵀ, L the lower triangle of scale_tril. L is finite, its diagonal in (0, inf).

Raises Invalid_argument if scale_tril's last two axes are not square or not of loc's length.

von_mises ~loc ~concentration is the distribution of angles on [-π, π] with density exp (κ cos (x - loc)) / (2π I₀(κ)), κ the concentration, in [0, inf): uniform at κ = 0, near a normal of variance 1 / κ about loc for a large κ. The log density is -2κ sin² ((x - loc) / 2) - log (2π i0e κ) (Nx.i0e), finite at every finite κ. loc is finite and may lie off [-π, π]: the density is periodic in it. Draws are Nx.Rng.von_mises's, Best and Fisher's rejection, moved by loc and wrapped onto [-π, π].

lkj_cholesky n ~concentration is the distribution of the Cholesky factor L of an n × n correlation matrix C = L Lᵀ whose density is proportional to det C^(η - 1), η the concentration in (0, inf) (Lewandowski, Kurowicka and Joe, Generating random correlation matrices based on vines and extended onion method, 2009): uniform over correlation matrices at η = 1, concentrated about the identity above. Over L's strict lower triangle the density is Π_(i = 2)^n L_ii^(n - i + 2η - 2) over their normalising constant, whose factors are formed from beta densities at their mode, which keep their digits at a large η. A matrix of concentrations' leading axes are batch axes, the value's last two axes the matrix. Draws follow the onion method; each off-diagonal C_ij has (C_ij + 1) / 2 beta of shapes η - 1 + n / 2 and η - 1 + n / 2. A matrix whose rows' squared norms are within 4 n ε of one counts as in the support.

Raises Invalid_argument if n < 2.

bernoulli ~logits is true with probability sigmoid logits. A logit may be infinite.

poisson ~rate is the count with probability rate^x exp (-rate) / x!, rate non-negative and finite.

neg_binomial ~mean ~dispersion is the count of mean mean and variance mean + mean² / dispersion, a Poisson whose rate is gamma with concentration dispersion. mean is non-negative and finite, dispersion in (0, inf).

binomial ~trials ~logits is the number of successes in trials independent trials, each a success with probability p = sigmoid logits: the count x in {0, ..., trials} with probability C(trials, x) p^x (1 - p)^(trials - x). trials is non-negative. A logit may be infinite: -inf puts all mass at 0, inf at trials. Draws are Nx.Rng.binomial's, within the limits it states.

categorical ~logits is the index k with probability softmax logits along the last axis. A logit may be -inf, a category of probability zero.

Raises Invalid_argument if logits is a scalar or its last axis is empty.

iid s d is s independent draws of d, its value of shape s followed by d's.

sorted n d is n independent draws of the scalar d sorted along a new last axis, with the exact density n! times the product of d's on nondecreasing vectors.

Raises Invalid_argument if d is not over scalars or n < 1.

transform b d is the distribution of b's image of a draw of d, with the density of d at the preimage less b's log-determinant there. A value outside b's image has density -inf.

truncated ?low ?high d is d conditioned on lying between low and high, elementwise, a missing bound leaving its side as d's: the density of d over its mass F(high) - F(low) there, -inf outside. The bounds are finite, broadcast against d's shape, and clip to d's support. The mass is a difference of upper tails where low lies above d's median and of lower tails otherwise, each held as logarithms, so the mass between bounds far in one tail keeps its digits; where the two tails agree within a factor e^(1/4) it is the density integrated over the window by eight-point Gauss-Legendre quadrature. The mass must be positive, a check like a parameter's. The log density is d's less the log mass: where both are large, their difference is within ε times their magnitude, as at bounds a million standard deviations out at float32. Draws and the quantile invert d at (1 - p) F(low) + p F(high), read from the smaller tail, within that inverse's precision, about ε min (F, 1 - F) / f in the value: a window narrower than that is not resolved. Its coordinates map onto the interval, or the half-line, between the bounds and d's support.

Raises Invalid_argument if neither bound is given, or if d is not normal, half_normal, lognormal, cauchy, half_cauchy, laplace, logistic, exponential, gamma, inverse_gamma, beta, uniform or iid of one: the families whose distribution function and its inverses nx computes.

mixture ~logits d draws one component k with probability softmax logits for the whole value, then the value from d_k, the slice of d at k on its leading axis: the log density is logsumexp_k (log_softmax logits_k + log_density d_k x). A component per point is iid [| n |] (mixture ~logits d). Its coordinates map onto the hull of its components' supports, read on the host.

Raises Invalid_argument if logits is not a vector or d's leading axis is not its length.

hmm ~initial ~transition n d is n observations of a hidden Markov chain over k states, its states summed out: state z_1 is j with probability softmax initial, state z_(t+1) is j with probability softmax transition_i after z_t = i, and observation t is a draw of d_(z_t), the slice of d at z_t on its leading axis. The value has shape [n] followed by d's components'. The log density is the forward algorithm's in logs (Rabiner, A Tutorial on Hidden Markov Models, 1989, §III-A), logsumexp_j α_n(j), its n - 1 products of k × k matrices multiplied pairwise in ⌈log₂ n⌉ rounds. membership gives each step's state's posterior from the forward and backward variables, scanned in 2 log₂ n rounds (Nx.associative_scan). A logit may be -inf, a state or a transition of probability zero. Its coordinates are those of a mixture of d.

Raises Invalid_argument if initial is not a vector of k, transition not a k × k matrix, d's leading axis not k, or n < 1; its eliminators raise on a value whose step axis is not of n.

log_density d x is Nx.sum (factors d x), a scalar.

factors d x is the log density of each independent unit of x: one per element for the elementwise families, per row for dirichlet and mvn, per matrix for lkj_cholesky, one for sorted, mixture and hmm, with iid's axes in front. x broadcasts against the parameters. A unit outside the support has factor -inf.

sample k d is a draw of d from the key k. It is reparameterised, differentiable in the parameters, wherever rune differentiates the Nx.Rng sampler it uses.

quantile d p is the value whose CDF is p, elementwise, read from the smaller tail: above 1/2 it inverts the survival at 1 - p, which is exact there. For gamma and inverse_gamma it is 0. at p = 0. and inf at p = 1.; for beta, Nx.betaincinv, 0. at p = 0. and 1. at p = 1..

Rune differentiates it in p and in d's parameters. For gamma, inverse_gamma and beta the first derivatives are finite wherever the exact ones are; inverse_gamma's higher derivatives can overflow where the quantile passes about 2^512.

Raises Invalid_argument for a family whose quantile nx does not compute: student_t, von_mises, the vector and matrix families and the combinators but iid and truncated.

coords d is the bijector onto d's support, from coordinates in which samplers and optimisers move: Bij.identity on the reals, Bij.exp on (0, inf), Bij.interval for beta, uniform and von_mises, Bij.simplex for dirichlet, Bij.cholesky_corr for lkj_cholesky, Bij.interval, Bij.greater or Bij.less for truncated, and for sorted d's bijector after Bij.ordered.

standardize d is the bijector whose coordinates are d's draws standardised: (x - loc) / scale for normal, student_t, cauchy, laplace and logistic, (log x - loc) / scale for lognormal, L⁻¹ (x - loc) for mvn, and the same for iid of these.

Raises Invalid_argument naming another family.

membership d x is the posterior probability of each value of d's discrete latent variables given x: of shape [k], each component's, for a mixture of k components; of shape [n; k], each step's state's, for an hmm of n steps over k states, from the forward and backward variables (γ_t ∝ α_t β_t, Rabiner 1989, eq. 27); and s followed by either for iid s of one.

Raises Invalid_argument if d is not one of those.

family d is d's family, such as "normal". iid keeps its argument's.

kind d is the type of d's values.

dtype d is the dtype of d's log densities, and of its values for a continuous d.

shape d is the shape of d's values.

placement d is Nx.beside of d's first parameter: where its bounds lie, and where sample places its key, so a draw lies where the parameters do whatever the key's placement. A transformation, a mixture or iid keeps its base's.

support d is the set of d's values. Bounds given by parameters are read on the host, as the interval holding every element's support. Inside a function Rune.jit compiles, reading one raises Invalid_argument, and inside a mapped function, reading one that depends on the function's arguments does. bounds reads nothing.

bounds d is the least and greatest value each element of d's values may take, of d's shape: -inf and inf where unbounded, the support's interval for an element of a vector. It reads nothing on the host, so it serves inside compiled and mapped functions.

pp ppf d formats d's family and its parameters' dtypes and shapes, as normal(loc: float64 [], scale: float64 [8]).

check ?unless what d is d whose parameters are checked now, and whose eliminators then check none: a caller that computes parameters names itself in the refusal. A parameter outside its domain raises Invalid_argument with the message what: family: param at [i] is v, not in domain. Where unless holds, a parameter is taken whatever its value: unless broadcasts against the parameters.

valid d is whether every element of every parameter of d is in its domain, a scalar over all of d's parameters: under Rune.vmap, one per lane. NaN is in no domain. It reads the checks check runs, raises nothing and ignores whether d was checked.