Links

Link functions connect model predictions to likelihood distributions, abstracting away NumPyro boilerplate for common output types.

Usage:

from blayers.layers import AdaptiveLayer
from blayers.links import gaussian_link

def model(x, y=None):
    mu = AdaptiveLayer()("mu", x)
    return gaussian_link(mu, y)

# HalfNormal sigma instead of Exponential
from functools import partial
import numpyro.distributions as dists
hn_gaussian = partial(gaussian_link, sigma_dist=dists.HalfNormal, sigma_kwargs={"scale": 1.0})

Available links:

• gaussian_link — Normal likelihood, configurable sigma prior

• lognormal_link — LogNormal likelihood, configurable sigma prior

• student_t_link — StudentT likelihood for robust regression (default df=4)

• gamma_link — Gamma likelihood (log link) for positive continuous data

• exponential_link — Exponential likelihood (log link) for positive / survival data

• logit_link — Bernoulli likelihood (binary)

• categorical_link — Categorical / softmax likelihood (multiclass)

• poisson_link — Poisson likelihood

• negative_binomial_link — NegativeBinomial2 likelihood, learned concentration

• ordinal_link — Ordinal (cumulative logit / proportional odds)

• zip_link — Zero-inflated Poisson

• zinb_link — Zero-inflated NegativeBinomial2 (overdispersed counts)

• beta_link — Beta regression for proportions in (0, 1)

blayers.links.gaussian_link(y_hat, y=None, *, obs_dist=<class 'numpyro.distributions.continuous.Normal'>, sigma_dist=<class 'numpyro.distributions.continuous.Exponential'>, sigma_kwargs=None, scale=None, untransformed_scale=None)

Gaussian likelihood with configurable sigma prior.

Default: sigma ~ Exponential(rate=1.0). Override via sigma_dist / sigma_kwargs. Pass a known scale or a raw untransformed_scale (transformed via softplus internally) to skip the sigma sample site.

Parameters:

• y_hat (Array) – Predicted mean, shape (n, 1) or (n,).

• y (Array | None) – Observed values, or None for prior predictive / inference.

• sigma_dist (Any) – Prior distribution class for sigma. Default Exponential.

• sigma_kwargs (dict[str, Any] | None) – Kwargs for sigma_dist. Default {"rate": 1.0}.

• scale (float | Array | None) – Known positive standard deviation. Scalar or broadcastable array.

• untransformed_scale (Array | None) – Unbounded array transformed via softplus internally.

• obs_dist (Any)

Returns:

Sample site "obs".

Return type:

Array

Example:

# Default: Exponential(1) prior on sigma
gaussian_link(mu, y)

# HalfNormal prior instead
from functools import partial
hn_link = partial(gaussian_link, sigma_dist=dists.HalfNormal, sigma_kwargs={"scale": 1.0})

# Known sigma (e.g. from XGBoost quantile regression)
gaussian_link(mu, y, scale=pred_std)

# Learned scale from a layer — softplus applied internally
raw = AdaptiveLayer()("log_scale", x)
gaussian_link(mu, y, untransformed_scale=raw)

blayers.links.lognormal_link(y_hat, y=None, *, obs_dist=<class 'numpyro.distributions.continuous.LogNormal'>, sigma_dist=<class 'numpyro.distributions.continuous.Exponential'>, sigma_kwargs=None, scale=None, untransformed_scale=None)

LogNormal likelihood with configurable sigma prior.

Default: sigma ~ Exponential(rate=1.0).

Parameters:

• y_hat (Array) – Log-scale predicted mean, shape (n, 1) or (n,).

• y (Array | None) – Observed positive values, or None.

• sigma_dist (Any) – Prior distribution class for sigma. Default Exponential.

• sigma_kwargs (dict[str, Any] | None) – Kwargs for sigma_dist. Default {"rate": 1.0}.

• scale (float | Array | None) – Known positive standard deviation.

• untransformed_scale (Array | None) – Unbounded array transformed via softplus internally.

• obs_dist (Any)

Returns:

Sample site "obs".

Return type:

Array

blayers.links.student_t_link(y_hat, y=None, *, obs_dist=functools.partial(<class 'numpyro.distributions.continuous.StudentT'>, df=4.0), sigma_dist=<class 'numpyro.distributions.continuous.Exponential'>, sigma_kwargs=None, scale=None, untransformed_scale=None)

StudentT likelihood for robust regression.

Heavier tails than Gaussian — large residuals are down-weighted rather than driving the fit. Default df=4 gives moderate robustness. Customise via functools.partial:

from functools import partial
cauchy_link = partial(student_t_link, obs_dist=partial(dists.StudentT, df=1.0))

Parameters:

• y_hat (Array) – Predicted location, shape (n, 1) or (n,).

• y (Array | None) – Observed values, or None.

• sigma_dist (Any) – Prior for scale. Default Exponential(rate=1.0).

• sigma_kwargs (dict[str, Any] | None) – Kwargs for sigma_dist.

• scale (float | Array | None) – Known positive scale.

• untransformed_scale (Array | None) – Unbounded scale transformed via softplus internally.

• obs_dist (Any)

Returns:

Sample site "obs".

Return type:

Array

blayers.links.gamma_link(y_hat, y=None, rate=1.0)

Gamma likelihood (log link) for positive continuous data.

Uses a mean parameterisation with a learned shape k:

\[\mu = \exp(\hat{y}), \quad k \sim \mathrm{Exponential}(\text{rate}), \quad y \sim \mathrm{Gamma}(k,\; k / \mu)\]

so E[y] = mu and Var[y] = mu^2 / k.

Parameters:

• y_hat (Array) – Log mean, shape (n, 1) or (n,).

• y (Array | None) – Observed positive values, or None.

• rate (float) – Rate of the Exponential prior on the shape k.

Returns:

Sample site "obs".

Return type:

Array

blayers.links.exponential_link(y_hat, y=None)

Exponential likelihood (log link) for positive continuous / survival data.

\[\mu = \exp(\hat{y}), \quad y \sim \mathrm{Exponential}(1 / \mu)\]

A single-parameter special case of gamma_link() (shape fixed at 1); the mean fully determines the variance (Var[y] = mu^2).

Parameters:

• y_hat (Array) – Log mean, shape (n, 1) or (n,).

• y (Array | None) – Observed positive values, or None.

Returns:

Sample site "obs".

Return type:

Array

blayers.links.logit_link(y_hat, y=None)

Bernoulli likelihood for binary classification.

Parameters:

• y_hat (Array) – Log-odds (logits), shape (n, 1) or (n,).

• y (Array | None) – Binary observations in {0, 1}, or None.

Returns:

Sample site "obs".

Return type:

Array

blayers.links.categorical_link(logits, y=None)

Categorical (softmax) likelihood for multiclass classification.

The multiclass generalisation of logit_link(). Produce one logit per class with a layer’s units argument (units = num_classes); the number of classes is read from the trailing dimension of logits.

\[P(Y = k \mid \text{logits}) = \mathrm{softmax}(\text{logits})_k\]

Parameters:

• logits (Array) – Unnormalised class scores of shape (n, num_classes) — e.g. AdaptiveLayer()("beta", x, units=K). A trailing singleton ((n, num_classes, 1)) is squeezed automatically.

• y (Array | None) – Integer class labels in {0, ..., num_classes - 1}, or None for prior predictive / inference.

Returns:

Sample site "obs" with integer values in {0, …, num_classes-1}.

Return type:

Array

Example:

from blayers.layers import AdaptiveLayer
from blayers.links import categorical_link

def model(x, y=None):
    logits = AdaptiveLayer()("beta", x, units=4)   # 4 classes
    return categorical_link(logits, y)

blayers.links.poisson_link(y_hat, y=None)

Poisson likelihood for count data.

Parameters:

• y_hat (Array) – Log rate, shape (n, 1) or (n,).

• y (Array | None) – Non-negative integer observations, or None.

Returns:

Sample site "obs".

Return type:

Array

blayers.links.negative_binomial_link(y_hat, y=None, rate=1.0)

NegativeBinomial2 likelihood for overdispersed count data.

Parameters:

• y_hat (Array) – Predicted mean, shape (n, 1) or (n,).

• y (Array | None) – Non-negative integer observations, or None.

• rate (float) – Rate parameter for the Exponential prior on concentration.

Returns:

Sample site "obs".

Return type:

Array

blayers.links.ordinal_link(mu, y=None, *, num_classes)

Cumulative logit (proportional odds) link for ordinal outcomes.

Models P(Y = k | μ) via:

\[P(Y \leq k \mid \mu) = \sigma(c_k - \mu)\]

Cutpoints are sampled with an ordered parameterisation: the first is free (Normal(0, 2)), subsequent ones add Exponential increments.

Parameters:

• mu (Array) – Linear predictor, shape (n, 1) or (n,).

• y (Array | None) – Integer observations in {0, 1, ..., num_classes - 1}, or None for prior predictive / inference.

• num_classes (int) – Number of ordinal categories (required).

Returns:

Sample site "obs" with integer values in {0, …, num_classes-1}.

Return type:

Array

blayers.links.zip_link(mu, y=None)

Zero-inflated Poisson link for count data with excess zeros.

Models a mixture: with probability π the outcome is exactly 0; with probability 1 - π the outcome follows Poisson(exp(μ)). π is a global scalar learned from data.

Parameters:

• mu (Array) – Log Poisson rate, shape (n, 1) or (n,).

• y (Array | None) – Non-negative integer observations, or None.

Returns:

Sample site "obs".

Return type:

Array

blayers.links.zinb_link(mu, y=None, rate=1.0)

Zero-inflated NegativeBinomial2 link for overdispersed counts with excess zeros.

The overdispersed counterpart of zip_link(): a mixture that emits an exact 0 with probability π, and otherwise a NegativeBinomial2 count with mean = exp(μ) and a learned concentration (so the non-zero part can be more dispersed than Poisson).

\[\text{gate} \sim \mathrm{Beta}(1, 10), \quad \phi \sim \mathrm{Exponential}(\text{rate}), \quad y \sim \mathrm{ZINB}(\exp(\mu),\; \phi,\; \text{gate})\]

Parameters:

• mu (Array) – Log mean, shape (n, 1) or (n,).

• y (Array | None) – Non-negative integer observations, or None.

• rate (float) – Rate of the Exponential prior on the concentration.

Returns:

Sample site "obs".

Return type:

Array

blayers.links.beta_link(mu, y=None)

Beta likelihood for proportional outcomes strictly in (0, 1).

Maps the linear predictor to a mean via sigmoid, then uses a learned global precision φ:

\[\bar{\mu} = \sigma(\mu), \quad y \sim Beta(\bar{\mu}\,\phi,\; (1 - \bar{\mu})\,\phi)\]

Parameters:

• mu (Array) – Logit of the mean proportion, shape (n, 1) or (n,).

• y (Array | None) – Observed proportions in (0, 1), or None.

Returns:

Sample site "obs".

Return type:

Array