Installation#

normix requires Python ≥ 3.12 and runs on CPU or GPU through JAX.

From PyPI#

pip install normix

Optional plotting helpers:

pip install "normix[plotting]"

The docs site header shows the version of the package used to build these pages. PyPI releases can lag the docs branch — if pip installs an older wheel than the header, prefer the from-source path below for the matching code.

From source (development)#

git clone https://github.com/xshi19/normix
cd normix
uv sync

uv sync installs the locked dependency set into a project virtual environment. Run anything in that environment with uv run, e.g. uv run python, uv run pytest, uv run jupyter lab.

Editable install with pip instead of uv:

pip install -e .

Float64 precision#

normix relies on double precision throughout — the Bessel evaluations, the GIG \(\eta \mapsto \theta\) solve, and log-density arithmetic all lose accuracy in float32. Enable it before importing normix:

import jax
jax.config.update("jax_enable_x64", True)
import normix

If float64 is not enabled, normix emits a warning on import.

Dependencies#

Package

Role

JAX

Array computation, autodiff, JIT, vmap

Equinox

Immutable pytree modules

JAXopt

L-BFGS/BFGS for the GIG \(\eta \mapsto \theta\) solve

NumPy

CPU array helpers alongside the JAX triad

SciPy

CPU solvers (trust-exact, L-BFGS-B); kve is a Bessel test oracle

Optional extras#

uv sync --extra docs --extra plotting

The docs extra builds the Sphinx site; plotting adds matplotlib helpers used by the tutorials.

Once installed, head to the Quickstart for a 30-second example, or Your first model, step by step for a guided walkthrough.