Skip to content

Model Fitting¤

glmax.fit(family, X, y) takes a family and observed data arrays and returns a fitted noun. The fitting strategy is an explicit fitter= argument — default IRLSFitter, or NewtonFitter for Fisher scoring Newton with backtracking line search — that can be swapped without changing anything else in the workflow.

Offsets and Exposure¤

offset is a fixed additive term in the linear predictor:

\[ \eta = X\beta + \mathrm{offset}. \]

For Poisson and Negative Binomial models with a log link, exposure is represented as an offset on the log scale:

fitted = glmax.fit(glmax.Poisson(), X, y, offset=jnp.log(exposure))

This models expected counts as \(\mu = \mathrm{exposure} \cdot \exp(X\beta)\).

Pass log exposure, not raw exposure

offset is added directly to \(\eta\). For exposure time, area, population at risk, sequencing depth, or similar denominators, pass jnp.log(exposure). Raw exposure belongs in the mean model only after taking the log.

glmax.fit(family: glmax.ExponentialDispersionFamily, X: ArrayLike, y: ArrayLike, *, offset: ArrayLike | None = None, init: glmax.Params | None = None, fitter: glmax.AbstractFitter = glmax.IRLSFitter(solver=lineax.Cholesky(), step_size=1.0, tol=0.001, max_iter=1000)) -> glmax.FittedGLM ¤

Fit a GLM to observed data and return a fitted noun.

This is the canonical fit grammar verb. It is @eqx.filter_jit-wrapped; the fitter strategy is treated as static structure under JIT. The returned value is a glmax.FittedGLM noun that binds the family and the full glmax.FitResult contract.

Differentiation is supported via a registered custom_jvp rule based on the Implicit Function Theorem: at convergence the score is zero, so the tangent for \(\hat\beta\) satisfies \(H \, d\hat\beta = -\partial_{\text{data}} \nabla_\beta \ell \cdot d(\text{data})\), where \(H = X^\top W X\) is the Fisher information at the converged fit.

Arguments:

  • family: glmax.ExponentialDispersionFamily instance.
  • X: covariate matrix, shape (n, p).
  • y: response vector, shape (n,).
  • offset: optional offset vector added to the linear predictor. For log-link count models with exposure, pass log(exposure), not raw exposure.
  • init: optional glmax.Params for warm-starting; None uses the family default.
  • fitter: glmax.AbstractFitter strategy. Defaults to glmax.IRLSFitter.

Returns:

glmax.FittedGLM noun binding the family and the glmax.FitResult contract.

Raises:

  • TypeError: if family, init, or fitter have wrong types.
  • ValueError: if array shapes or values are invalid.

glmax.Params(tuple) ¤

Canonical model parameter carrier.

A lightweight immutable container for the fitted coefficient vector, dispersion estimate, and optional family-specific auxiliary scalar. Used as the init argument to glmax.fit for warm-starting and forwarded inside glmax.FitResult and glmax.InferenceResult. The tuple stores \((\beta, \phi, a)\), where \(\beta\) is the regression coefficient vector, \(\phi\) is the GLM dispersion scalar, and \(a\) is optional family-specific auxiliary state.

Arguments:

  • beta: coefficient vector \(\beta\), inexact rank-1 array of shape (p,).
  • disp: dispersion scalar \(\phi\).
  • aux: optional auxiliary scalar \(a\).
aux class-attribute ¤

Alias for field number 2

beta class-attribute ¤

Alias for field number 0

disp class-attribute ¤

Alias for field number 1

glmax.FitResult ¤

Canonical fitter output contract.

Produced by every glmax.AbstractFitter strategy and consumed by glmax.FittedGLM, glmax.infer, and glmax.check. All fields are validated at construction time via __check_init__.

beta property ¤
__init__(self, params: glmax.Params, X: jax.Array, y: jax.Array, eta: jax.Array, mu: jax.Array, glm_wt: jax.Array, converged: jax.Array, num_iters: jax.Array, objective: jax.Array, objective_delta: jax.Array, score_residual: jax.Array) ¤

Construct a validated glmax.FitResult.

The fitted state stores \((\hat{\beta}, \hat{\eta}, \hat{\mu})\), where \(\hat{\beta}\) is the fitted coefficient vector, \(\hat{\eta}\) is the fitted linear predictor, and \(\hat{\mu}\) is the fitted mean on the response scale.

Arguments:

  • params: glmax.Params holding \((\hat{\beta}, \hat{\phi}, \hat{a})\).
  • X: covariate matrix \(X\), shape (n, p).
  • y: observed response vector \(y\), shape (n,).
  • eta: fitted linear predictor \(\hat{\eta}\), shape (n,).
  • mu: fitted mean response \(\hat{\mu}\), shape (n,).
  • glm_wt: GLM working weight vector \(w\), shape (n,).
  • converged: boolean scalar; True if the fitter converged within tolerance.
  • num_iters: integer scalar; number of iterations taken.
  • objective: final negative log-likelihood scalar.
  • objective_delta: change in objective on the last iteration.
  • score_residual: score-style residual, shape (n,).

glmax.FittedGLM ¤

Canonical fitted-model noun produced by fit(...).

Binds an glmax.ExponentialDispersionFamily with its glmax.FitResult and forwards the most commonly accessed fit artifacts as properties. Pass glmax.FittedGLM directly to glmax.infer and glmax.check.

Common artifacts are available as forwarding properties: params, beta, eta, mu, glm_wt, converged, num_iters, objective, objective_delta, score_residual.

X property ¤
beta property ¤
converged property ¤
eta property ¤
glm_wt property ¤
mu property ¤
num_iters property ¤
objective property ¤
objective_delta property ¤
params property ¤
score_residual property ¤
y property ¤
__init__(self, family: glmax.ExponentialDispersionFamily, result: glmax.FitResult) ¤

Construct a glmax.FittedGLM noun.

Arguments: