glmax

Grammar-first generalized linear modeling in JAX. JIT-compiled end-to-end and differentiable through the fitted parameters via the implicit function theorem.

Example

import jax.numpy as jnp
import glmax

X = jnp.array([[1.0, 0.5], [1.0, -0.3], [1.0, 1.2], [1.0, -0.8]])
y = jnp.array([2.0, 1.0, 4.0, 1.0])

fitted  = glmax.fit(glmax.Poisson(), X, y)
pred    = glmax.predict(fitted.family, fitted.params, X)
result  = glmax.infer(fitted)
diag    = glmax.check(fitted)

Four verbs — fit, predict, infer, and check — cover the full modeling workflow. Each takes explicit inputs and returns an explicit result. No hidden state is threaded between calls. predict returns response-scale fitted means, matching fitted.mu.

See the docs for the full API reference and guides.

Installation

For the latest release:

pip install glmax

For the latest repository state:

pip install git+https://github.com/mancusolab/glmax.git

Performance

Benchmarked against statsmodels on Poisson regression. Timing uses 10 steady-state runs after JIT warm-up. Statsmodels does not support GPU/TPU natively, and all reported runtimes are CPU-based, however the colab environment has different host CPUs depending on the runtime selected. As such we report the native CPU runtimes for statsmodels for each runtime environment.

See examples/benchmark_colab.ipynb for the full benchmark notebook.

Testing

pytest -p no:capture tests

Notes

glmax was developed by members of the Mancuso Lab with assistance from Claude Code and Codex, following the practices described in the scientific-software-playbook.

License

MIT