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

• logit_link — Bernoulli likelihood

• poisson_link — Poisson likelihood

• negative_binomial_link — NegativeBinomial2 likelihood, learned concentration

• ordinal_link — Ordinal (cumulative logit / proportional odds)

• zip_link — Zero-inflated Poisson

• 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 – Prior distribution class for sigma. Default Exponential.

• sigma_kwargs (dict | 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.

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 – Prior distribution class for sigma. Default Exponential.

• sigma_kwargs (dict | 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.

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.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=None)

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.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