Metadata-Version: 2.4
Name: glmax
Version: 0.1.1
Summary: Grammar-first generalized linear modeling in JAX.
Project-URL: Documentation, https://github.com/MancusoLab/glmax#readme
Project-URL: Issues, https://github.com/MancusoLab/glmax/issues
Project-URL: Source, https://github.com/MancusoLab/glmax
Author-email: Linda Serefin <lserafin@usc.edu>, Zixuan Zhang <zzhang39@usc.edu>, Nicholas Mancuso <nmancuso@usc.edu>
License-Expression: MIT
License-File: LICENSE
Keywords: glm,jax,machine-learning,regression,statistics
Classifier: Development Status :: 4 - Beta
Classifier: Intended Audience :: Developers
Classifier: Intended Audience :: Science/Research
Classifier: License :: OSI Approved :: MIT License
Classifier: Natural Language :: English
Classifier: Programming Language :: Python
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: Python :: 3.13
Classifier: Programming Language :: Python :: 3.14
Classifier: Programming Language :: Python :: Implementation :: CPython
Classifier: Programming Language :: Python :: Implementation :: PyPy
Classifier: Topic :: Scientific/Engineering
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Classifier: Topic :: Scientific/Engineering :: Information Analysis
Classifier: Topic :: Scientific/Engineering :: Mathematics
Requires-Python: >=3.11
Requires-Dist: equinox>=0.13
Requires-Dist: jax>=0.8.0
Requires-Dist: jaxlib>=0.8.0
Requires-Dist: jaxtyping
Requires-Dist: lineax>=0.1.0
Provides-Extra: docs
Requires-Dist: mkdocstrings-python>=1.12; extra == 'docs'
Requires-Dist: zensical>=0.0.37; extra == 'docs'
Description-Content-Type: text/markdown

# 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

```python
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](https://mancusolab.github.io/glmax) for the full API reference and guides.

## Installation

For the latest release:

```bash
pip install glmax
```

For the latest repository state:

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

## Performance

Benchmarked against [statsmodels](https://www.statsmodels.org/) 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.

| n | p | statsmodels (ms) | glmax (ms) | speedup | runtime |
|------:|----:|-----------------:|-----------:|--------:|---------|
| 500 | 10 | 4.32 | 0.92 | 4.7× | CPU |
| 2,000 | 20 | 277.76 | 4.14 | 67.1× | CPU |
| 10,000 | 50 | 1428.76 | 42.77 | 33.4× | CPU |
| 500 | 10 | 2.97 | 3.00 | 1.0× | T4 GPU |
| 2,000 | 20 | 13.94 | 4.38 | 3.2× | T4 GPU |
| 10,000 | 50 | 212.70 | 17.94 | 11.9× | T4 GPU |
| 500 | 10 | 1.90 | 2.89 | 0.7× | v5e-1 TPU |
| 2,000 | 20 | 8.46 | 8.80 | 1.0× | v5e-1 TPU |
| 10,000 | 50 | 1220.66 | 25.65 | 47.6× | v5e-1 TPU |

See [`examples/benchmark_colab.ipynb`](examples/benchmark_colab.ipynb) for the full benchmark notebook.

## Testing

```bash
pytest -p no:capture tests
```

## Notes

`glmax` was developed by members of the Mancuso Lab with assistance from
[Claude Code](https://claude.ai/code) and [Codex](https://chatgpt.com/codex), following the practices described in the
[scientific-software-playbook](https://github.com/mancusolab/scientific-software-playbook).

## License

MIT
