Families and links

A GLM is defined by its response family and link function. Pass a family instance as the first argument to glmax.fit:

import glmax

fitted = glmax.fit(glmax.Poisson(), X, y)               # log link by default
fitted = glmax.fit(glmax.Binomial(glmax.ProbitLink()), X, y)  # explicit link

The family determines how the linear predictor \(\eta = X\beta\) maps to the mean response \(\mu\), how the variance scales with \(\mu\), and how glmax.Params fields are interpreted:

• disp is the GLM dispersion \(\phi\). Gaussian and Gamma use it as EDM dispersion; Poisson, Binomial, and Negative Binomial canonicalize it to 1.0.
• aux carries optional family-specific state. Negative Binomial stores its overdispersion alpha in aux while canonical disp remains 1.0.

---

Exponential dispersion families

ExponentialDispersionFamily defines the common interface that fitting, inference, diagnostics, and prediction rely on. Concrete families implement this contract.

glmax.ExponentialDispersionFamily

glmax.ExponentialDispersionFamily

Abstract base for one-parameter exponential dispersion family distributions.

A GLM models the conditional mean \(\mu = \mathrm{E}(Y \mid X)\) via a link function \(g\) such that \(g(\mu) = \eta = X \beta\).

Subclasses implement the family-specific density, variance function, and split (disp, aux) handling, where disp stores the dispersion parameter \(\phi\) and aux stores optional family-specific state \(a\).

negloglikelihood(self, y: Array, eta: Array, disp: Shaped[jaxlib._jax.Array, ''] = 0.0, aux: Shaped[jaxlib._jax.Array, ''] | None = None) -> Array

Compute negative log-likelihood.

This returns the scalar objective \(-\log p(y \mid \eta, \phi, a)\), where \(y\) is the response vector, \(\eta\) is the linear predictor, \(\phi\) is the dispersion scalar, and \(a\) is optional auxiliary state.

Arguments:

• y: observed responses, shape (n,).
• eta: linear predictor, shape (n,).
• disp: dispersion parameter, scalar.
• aux: optional family-specific auxiliary scalar.

Returns:

Scalar negative log-likelihood.

variance(self, mu: Array, disp: Shaped[jaxlib._jax.Array, ''] = 0.0, aux: Shaped[jaxlib._jax.Array, ''] | None = None) -> Array

Variance function \(V(\mu)\).

This returns the family-specific variance expression evaluated at the mean response \(\mu\). The dispersion scalar \(\phi\) is passed via disp when the family uses it.

Arguments:

• mu: mean parameter, shape (n,).
• disp: dispersion parameter, scalar.
• aux: optional family-specific auxiliary scalar.

Returns:

Variance, shape (n,).

cdf(self, y: Array, mu: Array, disp: Shaped[jaxlib._jax.Array, ''] = 0.0, aux: Shaped[jaxlib._jax.Array, ''] | None = None) -> Array

Cumulative distribution function for the family.

This returns \(F(y_i \mid \mu_i, \phi, a)\) elementwise, where \(F\) is the fitted cumulative distribution function.

Arguments:

• y: observed responses, shape (n,).
• mu: fitted means, shape (n,).
• disp: dispersion parameter, scalar.
• aux: optional family-specific auxiliary scalar.

Returns:

CDF values \(F(y_i \mid \mu_i)\), shape (n,), values in [0, 1].

deviance_contribs(self, y: Array, mu: Array, disp: Shaped[jaxlib._jax.Array, ''] = 0.0, aux: Shaped[jaxlib._jax.Array, ''] | None = None) -> Array

Per-observation deviance contributions for the family.

The returned vector contains one deviance contribution \(d_i\) for each observation \(i\).

Arguments:

• y: observed responses, shape (n,).
• mu: fitted means, shape (n,).
• disp: dispersion parameter, scalar.
• aux: optional family-specific auxiliary scalar.

Returns:

Non-negative deviance contributions, shape (n,).

sample(self, key: Array, eta: Array, disp: Shaped[jaxlib._jax.Array, ''] = 0.0, aux: Shaped[jaxlib._jax.Array, ''] | None = None) -> Array

Draw samples from the family's distribution.

Sampling uses the parameterization implied by the link-transformed mean \(\mu = g^{-1}(\eta)\) together with the family's dispersion scalar \(\phi\) and auxiliary state \(a\) when present.

Arguments:

• key: JAX PRNGKey.
• eta: linear predictor, shape (n,).
• disp: dispersion parameter, scalar.
• aux: optional family-specific auxiliary scalar.

Returns:

Samples, shape (n,), same dtype as JAX float default.

calc_weight(self, eta: Array, disp: Shaped[jaxlib._jax.Array, ''] = 0.0, aux: Shaped[jaxlib._jax.Array, ''] | None = None) -> tuple[Array, Array, Array]

Compute IRLS weights.

Clips \(\mu\) to _bounds and variance to (tiny, inf) before computing the working weights. This returns \((\mu, g'(\mu), w)\) where \(w = 1 / \left(V(\mu) [g'(\mu)]^2\right)\). Here \(\mu\) is the mean response, \(V(\mu)\) is the family variance function, and \(g\) is the link function.

Arguments:

• eta: linear predictor, shape (n,).
• disp: dispersion parameter, scalar.
• aux: optional family-specific auxiliary scalar.

Returns:

Tuple (mu, g_deriv, weight) each of shape (n,), where g_deriv is the per-sample link derivative \(g'(\mu_i)\) and weight is the per-sample GLM working weight \(w_i = 1 / (V(\mu_i) [g'(\mu_i)]^2)\).

init_eta(self, y: Array) -> Array

update_nuisance(self, X: Array, y: Array, eta: Array, disp: Shaped[jaxlib._jax.Array, ''], step_size: Shaped[jaxlib._jax.Array, ''] = 1.0, aux: Shaped[jaxlib._jax.Array, ''] | None = None) -> tuple[Array, Array | None]

Apply one nuisance-parameter update step inside the IRLS loop.

Returns the updated (disp, aux) pair. The base implementation is a no-op identity: it returns (disp, aux) unchanged. Families that estimate a nuisance parameter override this method and update whichever slot they own: disp for families with free dispersion \(\phi\), and aux for structural-parameter families with auxiliary state \(a\). The base implementation simply returns the incoming nuisance state unchanged.

Arguments:

• X: design matrix, shape (n, p).
• y: observed responses, shape (n,).
• eta: linear predictor at current iteration, shape (n,).
• disp: current EDM dispersion scalar.
• step_size: IRLS step-size multiplier.
• aux: optional family-specific auxiliary scalar.

Returns:

Tuple (new_disp, new_aux).

init_nuisance(self) -> tuple[Array, Array | None]

Return the default (disp, aux) pair used to seed the IRLS loop.

This returns the initial nuisance state \((\phi_0, a_0)\).

Returns:

(jnp.array(1.0), None) for families without auxiliary state.

glmax.Gaussian(glmax.ExponentialDispersionFamily)

Gaussian (normal) exponential family with density

\[f(y \mid \mu, \phi) = \frac{1}{\sqrt{2\pi\phi}} \exp\!\left(-\frac{(y - \mu)^2}{2\phi}\right)\]

The mean is \(\mu \in \mathbb{R}\), the dispersion is \(\phi = \sigma^2 > 0\), and the variance function is \(V(\mu) = 1\), so \(\mathrm{Var}(Y \mid \mu) = \phi\).

The canonical link is the identity \(g(\mu) = \mu\). Gaussian stores dispersion in disp and ignores aux.

__init__(self, glink: glmax.AbstractLink = glmax.IdentityLink())

Construct a Gaussian family.

Arguments:

• glink: link function (default: IdentityLink()).

---

glmax.Gamma(glmax.ExponentialDispersionFamily)

Gamma exponential family with density

\[f(y \mid \mu, \phi) = \frac{y^{1/\phi - 1} \exp(-y / (\mu\phi))} {\Gamma(1/\phi)(\mu\phi)^{1/\phi}}\]

The mean is \(\mu > 0\) and the variance is \(\phi V(\mu)\) with \(V(\mu) = \mu^2\).

The canonical link for Gamma is InverseLink (\(g(\mu) = 1/\mu\)). Gamma uses disp as EDM dispersion and ignores aux.

__init__(self, glink: glmax.AbstractLink = glmax.InverseLink())

Construct a Gamma family.

Arguments:

• glink: link function (default: InverseLink()).

---

glmax.InverseGaussian(glmax.ExponentialDispersionFamily)

Inverse-Gaussian exponential family with density

\[f(y \mid \mu, \phi) = \left(\frac{1}{2\pi\phi y^3}\right)^{1/2} \exp\!\left(-\frac{(y - \mu)^2}{2\phi\mu^2 y}\right), \quad y > 0\]

The mean is \(\mu > 0\) and the variance is \(\phi V(\mu)\) with \(V(\mu) = \mu^3\).

The canonical link is \(g(\mu) = \mu^{-2}\), i.e. PowerLink(-2.0). Inverse Gaussian uses disp as EDM dispersion and ignores aux.

__init__(self, glink: glmax.AbstractLink = glmax.PowerLink(power=f32[](jax)))

Construct an InverseGaussian family.

Arguments:

• glink: link function (default: PowerLink(-2.0)).

---

glmax.Poisson(glmax.ExponentialDispersionFamily)

Poisson exponential family for count responses with PMF

\[P(Y = y \mid \mu) = \frac{e^{-\mu} \mu^y}{y!}, \quad y \in \{0, 1, 2, \ldots\}\]

The mean is \(\mu > 0\) and the variance function is \(V(\mu) = \mu\), so \(\mathrm{Var}(Y \mid \mu) = \mu\) with fixed dispersion \(\phi = 1\).

The canonical link is \(g(\mu) = \log(\mu)\). Poisson fixes disp = 1.0 and ignores aux.

__init__(self, glink: glmax.AbstractLink = glmax.LogLink())

Construct a Poisson family.

Arguments:

• glink: link function (default: LogLink()).

---

glmax.Binomial(glmax.ExponentialDispersionFamily)

Binomial exponential family for binary responses with PMF

\[P(Y = y \mid \mu) = \mu^y (1 - \mu)^{1 - y}, \quad y \in \{0, 1\}\]

The mean is \(\mu \in (0, 1)\) and the variance function is \(V(\mu) = \mu(1 - \mu)\), so \(\mathrm{Var}(Y \mid \mu) = \mu(1 - \mu)\) with fixed dispersion \(\phi = 1\).

The canonical link is the logit \(g(\mu) = \log(\mu / (1 - \mu))\). Binomial fixes disp = 1.0 and ignores aux.

__init__(self, glink: glmax.AbstractLink = glmax.LogitLink())

Construct a Binomial family.

Arguments:

• glink: link function (default: LogitLink()).

---

glmax.NegativeBinomial(glmax.ExponentialDispersionFamily)

Negative-binomial (NB-2) exponential family for overdispersed counts with PMF

\[P(Y = y \mid \mu, \alpha) = \binom{y + r - 1}{y} \left(\frac{r}{r + \mu}\right)^r \left(\frac{\mu}{r + \mu}\right)^y, \quad r = 1/\alpha, \; y \in \{0, 1, 2, \ldots\}\]

The mean is \(\mu > 0\), the overdispersion is \(\alpha > 0\), and the variance function is \(V(\mu) = \mu + \alpha\mu^2\) with fixed \(\phi = 1\). As \(\alpha \to 0\) the distribution converges to Poisson.

Negative Binomial fixes disp = 1.0 and stores overdispersion in aux.

__init__(self, glink: glmax.AbstractLink = glmax.LogLink())

Construct a Negative Binomial family.

Arguments:

• glink: link function (default: LogLink()).

---

Link functions

Links connect the mean response \(\mu\) to the linear predictor \(\eta\). The abstract link contract documents the forward link, inverse link, and their derivatives so the family layer and fitting kernels can work against one interface.

glmax.AbstractLink

glmax.AbstractLink

Abstract base for GLM link functions \(g: \mu \mapsto \eta\).

A link function connects the mean parameter \(\mu = \mathrm{E}(Y | X)\) to the linear predictor \(\eta = X \beta\) via \(\eta = g(\mu)\), where \(X\) is the design matrix and \(\beta\) is the coefficient vector.

Every concrete link exposes four related maps: \(g(\mu)\), its derivative \(g'(\mu)\), the inverse link \(g^{-1}(\eta)\), and the inverse-link derivative \((g^{-1})'(\eta)\).

__call__(self, mu: Array) -> Array

Compute \(g(\mu) = \eta\).

Here \(\mu\) is the mean-response vector and \(\eta\) is the linear predictor.

Arguments:

• mu: mean parameter, shape (n,), entries in the family's support.

Returns:

Linear predictor \(\eta\), shape (n,).

inverse(self, eta: Array) -> Array

Compute \(g^{-1}(\eta) = \mu\).

Here \(\eta\) is the linear predictor and \(\mu\) is the implied mean response.

Arguments:

• eta: linear predictor, shape (n,).

Returns:

Mean parameter \(\mu\), shape (n,).

deriv(self, mu: Array) -> Array

Compute \(g'(\mu)\).

The derivative is evaluated elementwise with respect to the mean response \(\mu\).

Arguments:

• mu: mean parameter, shape (n,).

Returns:

Link derivative, shape (n,).

inverse_deriv(self, eta: Array) -> Array

Compute \((g^{-1})'(\eta)\).

The derivative is evaluated elementwise with respect to the linear predictor \(\eta\).

Arguments:

• eta: linear predictor, shape (n,).

Returns:

Inverse-link derivative, shape (n,).

glmax.IdentityLink(glmax.AbstractLink)

Identity link \(g(\mu) = \mu\).

Canonical link for the Gaussian family.

Function	Derivative
\(g(\mu) = \mu\)	\(g'(\mu) = 1\)
\(g^{-1}(\eta) = \eta\)	\((g^{-1})'(\eta) = 1\)

__init__(self)

Initialize self. See help(type(self)) for accurate signature.

---

glmax.LogLink(glmax.AbstractLink)

Log link \(g(\mu) = \log(\mu)\).

Canonical link for the Poisson family.

Function	Derivative
\(g(\mu) = \log(\mu)\)	\(g'(\mu) = 1/\mu\)
\(g^{-1}(\eta) = e^\eta\)	\((g^{-1})'(\eta) = e^\eta\)

__init__(self)

Initialize self. See help(type(self)) for accurate signature.

---

glmax.LogitLink(glmax.AbstractLink)

Logit link \(g(\mu) = \log(\mu / (1 - \mu))\).

Canonical link for the Binomial family. Here \(\sigma\) denotes the logistic sigmoid.

Function	Derivative
\(g(\mu) = \log\!\left(\frac{\mu}{1-\mu}\right)\)	\(g'(\mu) = \frac{1}{\mu(1-\mu)}\)
\(g^{-1}(\eta) = \sigma(\eta)\)	\((g^{-1})'(\eta) = \sigma(\eta)(1-\sigma(\eta))\)

__init__(self)

Initialize self. See help(type(self)) for accurate signature.

---

glmax.InverseLink(glmax.AbstractLink)

Reciprocal link \(g(\mu) = 1 / \mu\).

Canonical link for the Gamma family.

Function	Derivative
\(g(\mu) = 1/\mu\)	\(g'(\mu) = -1/\mu^2\)
\(g^{-1}(\eta) = 1/\eta\)	\((g^{-1})'(\eta) = -1/\eta^2\)

__init__(self)

Initialize self. See help(type(self)) for accurate signature.

---

glmax.PowerLink(glmax.AbstractLink)

Power link \(g(\mu) = \mu^p\).

Here \(p\) is the power exponent. Derivatives are computed via autodiff.

Function	Derivative
\(g(\mu) = \mu^p\)	\(g'(\mu) = p\mu^{p-1}\)
\(g^{-1}(\eta) = \eta^{1/p}\)	\((g^{-1})'(\eta) = \eta^{1/p-1}/p\)

__init__(self, power: float = 1.0)

Construct a power link.

Arguments:

• power: exponent \(p\). Default 1.0 (identity link).

---

glmax.NBLink(glmax.AbstractLink)

Negative-binomial link \(g(\mu) = \log(\alpha \mu / (1 + \alpha \mu))\).

Canonical link for the Negative Binomial family. Here \(\alpha > 0\) is the overdispersion parameter. The inverse-link derivative is computed via autodiff.

Function	Derivative
\(g(\mu) = \log\!\left(\frac{\alpha\mu}{1+\alpha\mu}\right)\)	\(g'(\mu) = \frac{1}{\mu(1+\alpha\mu)}\)
\(g^{-1}(\eta) = \frac{1}{\alpha\,\mathrm{expm1}(-\eta)}\)	\((g^{-1})'(\eta)\) via autodiff

__init__(self, alpha: float = 1.0)

Construct a Negative Binomial link.

Arguments:

• alpha: overdispersion parameter \(\alpha > 0\). Default 1.0.

---

glmax.ProbitLink(glmax.AbstractLink)

Probit link \(g(\mu) = \Phi^{-1}(\mu)\).

Second most common link for the Binomial family, widely used in bioassay and econometrics. Here \(\Phi\) is the standard normal CDF and \(\phi\) is the standard normal PDF. Derivatives are computed via autodiff.

Function	Derivative
\(g(\mu) = \Phi^{-1}(\mu)\)	\(g'(\mu) = 1/\phi(\Phi^{-1}(\mu))\)
\(g^{-1}(\eta) = \Phi(\eta)\)	\((g^{-1})'(\eta) = \phi(\eta)\)

__init__(self)

Initialize self. See help(type(self)) for accurate signature.

---

glmax.CLogLogLink(glmax.AbstractLink)

Complementary log-log link \(g(\mu) = \log(-\log(1 - \mu))\).

Canonical link for interval-censored survival models and log-Weibull regression. Asymmetric: probability approaches 0 slower than it approaches 1. Derivatives are computed via autodiff.

Function	Derivative
\(g(\mu) = \log(-\log(1-\mu))\)	\(g'(\mu) = -\frac{1}{(1-\mu)\log(1-\mu)}\)
\(g^{-1}(\eta) = 1 - e^{-e^\eta}\)	\((g^{-1})'(\eta) = e^{\eta - e^\eta}\)

__init__(self)

Initialize self. See help(type(self)) for accurate signature.

---

glmax.LogLogLink(glmax.AbstractLink)

Log-log link \(g(\mu) = -\log(-\log(\mu))\).

Mirror of CLogLog: asymmetric in the opposite direction, approaching 0 faster than it approaches 1. Derivatives are computed via autodiff.

Function	Derivative
\(g(\mu) = -\log(-\log(\mu))\)	\(g'(\mu) = -\frac{1}{\mu\log(\mu)}\)
\(g^{-1}(\eta) = e^{-e^{-\eta}}\)	\((g^{-1})'(\eta) = e^{-\eta - e^{-\eta}}\)

__init__(self)

Initialize self. See help(type(self)) for accurate signature.

---

glmax.SqrtLink(glmax.AbstractLink)

Square-root link \(g(\mu) = \sqrt{\mu}\).

Variance-stabilising link for the Poisson family: approximately stabilises variance for count data.

Function	Derivative
\(g(\mu) = \sqrt{\mu}\)	\(g'(\mu) = \frac{1}{2\sqrt{\mu}}\)
\(g^{-1}(\eta) = \eta^2\)	\((g^{-1})'(\eta) = 2\eta\)

__init__(self)

Initialize self. See help(type(self)) for accurate signature.

---

glmax.CauchitLink(glmax.AbstractLink)

Cauchit link \(g(\mu) = \tan(\pi(\mu - 1/2))\).

Heavy-tailed alternative to the probit link for the Binomial family. Robust to extreme observations near 0 or 1: tails decay as \(1/\eta^2\) rather than exponentially. Derivatives are computed via autodiff.

Function	Derivative
\(g(\mu) = \tan(\pi(\mu - \tfrac{1}{2}))\)	\(g'(\mu) = \frac{\pi}{\cos^2(\pi(\mu - \tfrac{1}{2}))}\)
\(g^{-1}(\eta) = \tfrac{1}{2} + \frac{\arctan(\eta)}{\pi}\)	\((g^{-1})'(\eta) = \frac{1}{\pi(1+\eta^2)}\)

__init__(self)

Initialize self. See help(type(self)) for accurate signature.