augax 0.3.0__tar.gz

augax-0.3.0/.github/release.yml

@@ -0,0 +1,19 @@
+changelog:
+  exclude:
+    labels:
+      - ignore-for-release
+    authors:
+      - dependabot
+  categories:
+    - title: Bug Fixes
+      labels:
+        - bug
+    - title: Enhancements
+      labels:
+        - enhancement
+    - title: Maintenance
+      labels:
+        - maintenance
+    - title: Other Changes
+      labels:
+        - "*"

augax-0.3.0/.github/workflows/build-wheels.yml

@@ -0,0 +1,170 @@
+name: Build and publish to PyPI
+
+on:
+  push:
+    tags:
+      - "v*" # v1.2.3 → PyPI, v1.2.3a1/b1/rc1 → TestPyPI
+  workflow_dispatch:
+
+jobs:
+  build_wheels:
+    name: Build wheels on ${{ matrix.os }}
+    permissions:
+      id-token: write # MANDATORY for trusted publishing to PyPI
+      contents: write # MANDATORY for the GitHub release action
+    strategy:
+      fail-fast: false
+      matrix:
+        include:
+          # Linux x86_64 — manylinux via cibuildwheel's default container
+          - os: ubuntu-latest
+            arch: x86_64
+          # Linux aarch64 — emulated via QEMU
+          - os: ubuntu-latest
+            arch: aarch64
+          # macOS Apple Silicon (arm64) -- pinned rather than macos-latest,
+          # since macos-latest now resolves to macOS 26 (Tahoe) and
+          # Homebrew's bottles for that host embed a minimum target of 26.0,
+          # which would make the wheel unusable on any older macOS. Bump
+          # this pin (and the target below) deliberately, not automatically.
+          - os: macos-15
+            arch: arm64
+            macosx_deployment_target: "15.0"
+          # Windows x86_64 -- pinned rather than windows-latest, since
+          # windows-latest now resolves to a windows-2025-vs2026 preview
+          # image whose bundled flang compiler fails to compile reference
+          # LAPACK's Fortran source (a compiler bug, not fixable here).
+          # Bump this pin deliberately once VS2026/flang has matured.
+          - os: windows-2022
+            arch: AMD64
+    runs-on: ${{ matrix.os }}
+    env:
+      CIBW_BUILD: "cp311-* cp312-* cp313-* cp314-*"
+      # jaxlib (needed at build time for XLA FFI headers) has no musllinux
+      # wheels published at all -- only manylinux, macOS, and Windows. That's
+      # an upstream gap, not something fixable from this workflow, so skip
+      # musllinux targets entirely rather than let every build fail on it.
+      CIBW_SKIP: "*-musllinux*"
+      CIBW_ARCHS: ${{ matrix.arch }}
+      CIBW_BUILD_VERBOSITY: 1
+      MACOSX_DEPLOYMENT_TARGET: ${{ matrix.macosx_deployment_target }}
+
+      # ── Linux ──────────────────────────────────────────────────────────────
+      # The default manylinux2014 image is too old: jaxlib and scipy no
+      # longer publish wheels for it, which forced pip to backtrack to
+      # ancient jax versions and eventually build scipy from source (which
+      # then failed on a missing OpenBLAS). Use manylinux_2_28 instead, which
+      # is current enough for modern prebuilt wheels of everything involved.
+      CIBW_MANYLINUX_X86_64_IMAGE: manylinux_2_28
+      CIBW_MANYLINUX_AARCH64_IMAGE: manylinux_2_28
+      # augax is self-contained C++ (it links no SuiteSparse and no OpenMP), so
+      # no platform needs system libraries.  jax comes from pyproject's
+      # [build-system] pin -- the build is isolated, so nothing is
+      # pip-installed here.  cibuildwheel v4 repairs Windows wheels with
+      # delvewheel by default.
+
+      # ── Every platform: test the repaired wheel before it can publish ─────
+      # The reference-sampler tests need polyagamma, which has no macOS arm64 or
+      # cp314 wheels; they run in unittests.yml.  These checks run everywhere.
+      CIBW_TEST_REQUIRES: pytest
+      CIBW_TEST_COMMAND: >
+        python {project}/ci/check_wheel.py &&
+        pytest -q {project}/tests/test_batching.py
+
+    steps:
+      - uses: actions/checkout@v6
+        with:
+          fetch-depth: 0  # setuptools-scm needs tags
+
+      # Enable QEMU for Linux aarch64 cross-compilation.
+      - name: Set up QEMU (Linux aarch64 only)
+        if: matrix.os == 'ubuntu-latest' && matrix.arch == 'aarch64'
+        uses: docker/setup-qemu-action@v3
+        with:
+          platforms: arm64
+
+      - name: Build wheels
+        # v2.x predates CPython 3.14, so the cp314-* selector above matched
+        # nothing and 3.14 wheels were silently never built.
+        uses: pypa/cibuildwheel@v4.2.1
+        with:
+          output-dir: wheelhouse
+
+      - name: Upload wheel artifacts
+        uses: actions/upload-artifact@v4
+        with:
+          name: wheels-${{ matrix.os }}-${{ matrix.arch }}
+          path: wheelhouse/*.whl
+
+  build_sdist:
+    name: Build source distribution
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          fetch-depth: 0  # setuptools-scm needs tags
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: "3.12"
+      - name: Build sdist
+        run: |
+          pip install build
+          python -m build --sdist --outdir wheelhouse
+      - name: Upload sdist artifact
+        uses: actions/upload-artifact@v4
+        with:
+          name: sdist
+          path: wheelhouse/*.tar.gz
+
+  publish:
+    name: Publish to PyPI and create release
+    needs: [build_wheels, build_sdist]
+    runs-on: ubuntu-latest
+    permissions:
+      id-token: write
+      contents: write
+    # MANDATORY for trusted publishing: the environment name here MUST match the
+    # one configured in the PyPI/TestPyPI trusted-publisher UI. Without it, the
+    # OIDC token exchange fails with `invalid-publisher` / `environment: MISSING`.
+    environment: pypi
+    # Runs on any tag. test-* tags build+artifact only (no publish step fires).
+    if: startsWith(github.ref, 'refs/tags/')
+    steps:
+      - name: Download all artifacts
+        uses: actions/download-artifact@v4
+        with:
+          path: dist
+          merge-multiple: true
+
+      # v1.2.3a1 / v1.2.3b2 / v1.2.3rc1 → TestPyPI
+      - name: Publish pre-release to TestPyPI
+        if: |
+          startsWith(github.ref_name, 'v') &&
+          (contains(github.ref_name, 'a') || contains(github.ref_name, 'b') || contains(github.ref_name, 'rc'))
+        uses: pypa/gh-action-pypi-publish@release/v1
+        with:
+          repository-url: https://test.pypi.org/legacy/
+          packages-dir: dist/
+
+      # v1.2.3 (clean) → real PyPI
+      - name: Publish release to PyPI
+        if: |
+          startsWith(github.ref_name, 'v') &&
+          !contains(github.ref_name, 'a') &&
+          !contains(github.ref_name, 'b') &&
+          !contains(github.ref_name, 'rc')
+        uses: pypa/gh-action-pypi-publish@release/v1
+        with:
+          packages-dir: dist/
+
+      # GitHub Release for any v* tag; marked pre-release when applicable
+      - name: Create GitHub Release
+        if: startsWith(github.ref_name, 'v')
+        uses: softprops/action-gh-release@v2
+        with:
+          prerelease: ${{ contains(github.ref_name, 'a') || contains(github.ref_name, 'b') || contains(github.ref_name, 'rc') }}
+          files: dist/*
+          generate_release_notes: true
+        env:
+          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}

augax-0.3.0/.github/workflows/build_docs.yml

@@ -0,0 +1,65 @@
+ name: Build Docs
+
+ on:
+   push:
+     # Sequence of patterns matched against refs/tags
+     tags:
+       - 'v*' # Push events to matching v*, i.e. v1.0, v20.15.10
+   workflow_dispatch:
+     inputs:
+       version:
+         description: Manual Doc Build
+         default: test
+         required: false
+ permissions:
+   contents: write   # required to push to gh-pages
+ jobs:
+   docs:
+     name: build & push docs
+     runs-on: ${{ matrix.os }}
+     timeout-minutes: 90
+     strategy:
+       matrix:
+         os: ['ubuntu-latest']
+         environment-file: [docs/docs_env.yml]
+         experimental: [false]
+     defaults:
+       run:
+         shell: bash -l {0}
+
+     steps:
+       - name: checkout repo
+         uses: actions/checkout@v4
+         with:
+           fetch-depth: 0  # setuptools-scm needs tags
+
+       - name: setup micromamba
+         uses: mamba-org/setup-micromamba@v2
+         with:
+           environment-file: ${{ matrix.environment-file }}
+           micromamba-version: 'latest'
+
+       - name: make docs
+         run: |
+           pip install -e . --no-deps --force-reinstall
+           cd docs; make html
+
+       - name: commit docs
+         run: |
+           git clone https://github.com/knaaptime/augax.git --branch gh-pages --single-branch gh-pages
+           cp -r docs/build/html/* gh-pages/
+           cd gh-pages
+           git config --local user.email "action@github.com"
+           git config --local user.name "GitHub Action"
+           git add .
+           git commit -m "Update documentation" -a || true
+           # The above command will fail if no changes were present,
+           # so we ignore the return code.
+
+       - name: push to gh-pages
+         uses: ad-m/github-push-action@master
+         with:
+            branch: gh-pages
+            directory: gh-pages
+            github_token: ${{ secrets.GITHUB_TOKEN }}
+            force: true

augax-0.3.0/.github/workflows/unittests.yml

@@ -0,0 +1,75 @@
+name: Tests
+
+on:
+  push:
+    branches: [main]
+  pull_request:
+    branches:
+      - "*"
+  schedule:
+    - cron: "0 0 * * 1,4"
+
+jobs:
+  Test:
+    name: ${{ matrix.os }}, ${{ matrix.environment-file }}
+    runs-on: ${{ matrix.os }}
+    strategy:
+      fail-fast: false
+      matrix:
+        os: [ubuntu-latest]
+        environment-file:
+          - ci/311.yml
+          - ci/312.yml
+          - ci/313.yml
+          - ci/314.yml
+        include:
+          - environment-file: ci/313.yml
+            os: macos-latest
+          - environment-file: ci/313.yml
+            os: windows-latest
+    defaults:
+      run:
+        shell: bash -l {0}
+
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          fetch-depth: 0  # setuptools-scm needs tags
+
+      - name: setup micromamba
+        uses: mamba-org/setup-micromamba@v1
+        with:
+          environment-file: ${{ matrix.environment-file }}
+          micromamba-version: 'latest'
+
+      - name: Install package
+        run: pip install -e .  --no-build-isolation --no-deps
+
+      - name: Test package
+        run: |
+          pytest --cov src/augax \
+           -v \
+           -r a \
+           -n auto \
+           --color yes \
+           --cov src \
+           --cov-append \
+           --cov-report term-missing \
+           --cov-report xml \
+           tests
+      - name: codecov
+        uses: codecov/codecov-action@v5
+        with:
+          token: ${{ secrets.CODECOV_TOKEN }}
+          file: ./coverage.xml
+          name: cholax-codecov
+
+      - name: Generate and publish the report
+        if: |
+          failure()
+          && steps.status.outcome == 'failure'
+          && github.event_name == 'schedule'
+          && github.repository_owner == 'knaaptime'
+        uses: xarray-contrib/issue-from-pytest-log@v1
+        with:
+          log-path: pytest-log.jsonl

augax-0.3.0/.gitignore

@@ -0,0 +1,107 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+# C extensions
+*.so
+notebooks/cache
+**.DS_Store
+
+# Distribution / packaging
+.Python
+env/
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*,cover
+.hypothesis/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# pyenv
+.python-version
+
+# celery beat schedule file
+celerybeat-schedule
+
+# dotenv
+.env
+
+# virtualenv
+.venv
+venv/
+ENV/
+
+# Spyder project settings
+.spyderproject
+
+# Rope project settings
+.ropeproject
+
+# Pycharm
+.idea
+
+# VSCode
+.vscode
+# pixi environments
+.pixi
+*.egg-info
+
+.claude/
+
+**.ruff_cache/
+
+**.pytest_cache

augax-0.3.0/CMakeLists.txt

@@ -0,0 +1,28 @@
+cmake_minimum_required(VERSION 3.18)
+project(augax LANGUAGES CXX)
+
+set(CMAKE_CXX_STANDARD 17)
+set(CMAKE_CXX_STANDARD_REQUIRED ON)
+
+find_package(Python COMPONENTS Interpreter Development.Module REQUIRED)
+
+# XLA FFI headers ship with jaxlib (header-only, nothing to link).
+execute_process(
+    COMMAND "${Python_EXECUTABLE}" -c "import jax.ffi; print(jax.ffi.include_dir())"
+    OUTPUT_VARIABLE XLA_FFI_INCLUDE_DIR
+    OUTPUT_STRIP_TRAILING_WHITESPACE
+    COMMAND_ERROR_IS_FATAL ANY)
+
+# Locate nanobind via its Python package.
+execute_process(
+    COMMAND "${Python_EXECUTABLE}" -m nanobind --cmake_dir
+    OUTPUT_VARIABLE NANOBIND_CMAKE_DIR
+    OUTPUT_STRIP_TRAILING_WHITESPACE
+    COMMAND_ERROR_IS_FATAL ANY)
+list(APPEND CMAKE_PREFIX_PATH "${NANOBIND_CMAKE_DIR}")
+find_package(nanobind REQUIRED CONFIG)
+
+nanobind_add_module(augax_cpp NOMINSIZE src/cpp/augax.cpp)
+target_include_directories(augax_cpp PRIVATE ${XLA_FFI_INCLUDE_DIR})
+
+install(TARGETS augax_cpp LIBRARY DESTINATION .)

augax-0.3.0/LICENSE.txt

@@ -0,0 +1,13 @@
+BSD 3-Clause License
+
+Copyright 2026 elijah knaap
+
+Redistribution and use in source and binary forms, with or without modification, are permitted provided that the following conditions are met:
+
+1. Redistributions of source code must retain the above copyright notice, this list of conditions and the following disclaimer.
+
+2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the following disclaimer in the documentation and/or other materials provided with the distribution.
+
+3. Neither the name of the copyright holder nor the names of its contributors may be used to endorse or promote products derived from this software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

augax-0.3.0/PKG-INFO

@@ -0,0 +1,105 @@
+Metadata-Version: 2.2
+Name: augax
+Version: 0.3.0
+Summary: JAX-native on-device augmentation draws (Pólya-Gamma) as XLA FFI custom calls
+License: BSD-3-Clause
+Requires-Python: >=3.11
+Requires-Dist: jax>=0.9.0
+Requires-Dist: numpy
+Provides-Extra: test
+Requires-Dist: pytest; extra == "test"
+Requires-Dist: scipy; extra == "test"
+Description-Content-Type: text/markdown
+
+# augax
+
+On-device draws for **conditionally Gaussian augmentation** in JAX, each an XLA
+FFI custom call: **Pólya-Gamma** sampling (`pg_sample`) and the **Kolmogorov
+scale** of the logistic's normal scale mixture (`ks_sample`).
+
+`augax.pg_sample(h, z, key)` draws `omega ~ PG(h, z)` from *inside* a JIT'd /
+`lax.scan`-ed / `pmap`-ed program with **no host round-trip** — the native
+replacement for `jax.pure_callback(random_polyagamma, ...)`, which pays a
+device↔host round-trip on every Pólya-Gamma Gibbs sweep and serializes under
+`jax.pmap`.
+
+```python
+import jax, augax
+jax.config.update("jax_enable_x64", True)
+
+@jax.jit
+def omega_update(h, z, key):
+    return augax.pg_sample(h, z, key)
+```
+
+## Method
+
+`pg_sample` follows the hybrid dispatch of
+[`zoj613/polyagamma`](https://github.com/zoj613/polyagamma) (`method=None`),
+changed where that scheme is measurably biased:
+
+- **`h > 50`**: a normal with the exact Pólya-Gamma mean and variance.
+- **`8 <= h <= 50`**: the **saddlepoint** rejection sampler (Windle, Polson &
+  Scott 2014, arXiv:1405.0506), O(1) per accepted draw.
+- **`h = 1`** (Bernoulli/logit), **or integer `h <= 4` with `|z| <= 1`**: the
+  exact **Devroye** sampler for `PG(1, z)` (Polson, Scott & Windle 2013), summed
+  `h` times.
+- **Otherwise, `1 < h < 8`** — e.g. Negative-Binomial `h = y + alpha`: the
+  **alternate** sampler (Windle, Polson & Scott 2014, Section 4).
+- **`h < 1`**: the Gamma-series representation — 20 terms plus a remainder
+  matched in mean and variance — exact in both moments.
+
+Each departure from `polyagamma` 2.0.2 removes a bias measured with 2M draws
+against the exact mean and variance:
+
+- **Alternate sampler, `1 < h < 4`, small `|z|`**: `polyagamma` runs +0.4–0.9% in
+  the mean and +1–5% in the variance. Its right-hand envelope uses the constant
+  `(π/2)^(h/2)`, where both its Gamma proposal and the density's tail carry
+  `(π/2)^h`; augax uses the latter.
+- **`h < 1`**: the alternate method assumes `h >= 1`; below it `polyagamma` runs
+  −0.4–0.8% in the mean. augax uses the Gamma series.
+- **`4 < h < 8` with `|z| <= 4`**: `polyagamma` uses the saddlepoint
+  approximation, 1–2% low in the variance. augax uses the exact alternate
+  sampler.
+
+The saddlepoint and alternate samplers are ported from `polyagamma`
+(BSD-3-Clause). Every regime is tested against the closed-form mean and
+variance.
+
+The RNG is a per-call `xoshiro256**` seeded (via `splitmix64`) from the JAX PRNG
+key, so draws are reproducible and each call owns its stream — safe when several
+threads call at once. Normals and exponentials come from ziggurat samplers.
+
+An array longer than 65,536 elements is split into blocks of that size, each
+with its own stream, and the blocks are drawn in parallel on every hardware
+thread (set `AUGAX_NUM_THREADS` to cap it). The blocks depend only on the array
+length, so a draw is the same on any machine and any thread count. On a large
+Negative-Binomial model (8.1M cells) this takes the draw from 3.6 s to 0.33 s on
+16 threads.
+
+## The Kolmogorov scale
+
+A standard logistic disturbance is a normal scale mixture,
+`eps | lam ~ N(0, lam^2)` with `lam = 2K` and `K` Kolmogorov distributed
+(Stefanski 1991). `ks_sample(u, key)` draws `lam` given `eps = u`, which makes a
+logistic latent variable conditionally Gaussian: the propagated disturbance of a
+latent SAR logit, for one. The draw is exact: `1/lam^2` is proposed from an
+inverse Gaussian, `IG(1/|u|, 1)` (a chi-square at `|u| < 1e-4`), and accepted
+with probability `S(v) = f(v) e^{v/2} <= 1`, `f` the density of `(2K)^2`,
+evaluated lazily by the squeeze of an alternating series — the Kolmogorov
+series above `v = 4/3` and its Jacobi transform below, where the first's
+partial sums are not bounds (Holmes & Held 2006). It is checked against the
+identity `E[exp(-t^2 lam^2 / 2)] = pi t / sinh(pi t)` and against the
+posterior moments by quadrature. `ks_sample_numpy(u, rng)` is the same C++
+sampler for host callers.
+
+## Status
+
+- ✅ All `h > 0`: the logit (`h = 1`), Negative-Binomial (`h = y + alpha`) and
+  large-shape regimes.
+- Dispatch is automatic: the sampler is selected from `(h, z)` with no API
+  change.
+- ✅ `ks_sample`: exact for every `u`, on the device and on the host.
+
+CPU-only (the sampler is cheap scalar work; the point is to keep it on-device
+next to the rest of a JIT'd Gibbs sweep).

augax-0.3.0/README.md

@@ -0,0 +1,92 @@
+# augax
+
+On-device draws for **conditionally Gaussian augmentation** in JAX, each an XLA
+FFI custom call: **Pólya-Gamma** sampling (`pg_sample`) and the **Kolmogorov
+scale** of the logistic's normal scale mixture (`ks_sample`).
+
+`augax.pg_sample(h, z, key)` draws `omega ~ PG(h, z)` from *inside* a JIT'd /
+`lax.scan`-ed / `pmap`-ed program with **no host round-trip** — the native
+replacement for `jax.pure_callback(random_polyagamma, ...)`, which pays a
+device↔host round-trip on every Pólya-Gamma Gibbs sweep and serializes under
+`jax.pmap`.
+
+```python
+import jax, augax
+jax.config.update("jax_enable_x64", True)
+
+@jax.jit
+def omega_update(h, z, key):
+    return augax.pg_sample(h, z, key)
+```
+
+## Method
+
+`pg_sample` follows the hybrid dispatch of
+[`zoj613/polyagamma`](https://github.com/zoj613/polyagamma) (`method=None`),
+changed where that scheme is measurably biased:
+
+- **`h > 50`**: a normal with the exact Pólya-Gamma mean and variance.
+- **`8 <= h <= 50`**: the **saddlepoint** rejection sampler (Windle, Polson &
+  Scott 2014, arXiv:1405.0506), O(1) per accepted draw.
+- **`h = 1`** (Bernoulli/logit), **or integer `h <= 4` with `|z| <= 1`**: the
+  exact **Devroye** sampler for `PG(1, z)` (Polson, Scott & Windle 2013), summed
+  `h` times.
+- **Otherwise, `1 < h < 8`** — e.g. Negative-Binomial `h = y + alpha`: the
+  **alternate** sampler (Windle, Polson & Scott 2014, Section 4).
+- **`h < 1`**: the Gamma-series representation — 20 terms plus a remainder
+  matched in mean and variance — exact in both moments.
+
+Each departure from `polyagamma` 2.0.2 removes a bias measured with 2M draws
+against the exact mean and variance:
+
+- **Alternate sampler, `1 < h < 4`, small `|z|`**: `polyagamma` runs +0.4–0.9% in
+  the mean and +1–5% in the variance. Its right-hand envelope uses the constant
+  `(π/2)^(h/2)`, where both its Gamma proposal and the density's tail carry
+  `(π/2)^h`; augax uses the latter.
+- **`h < 1`**: the alternate method assumes `h >= 1`; below it `polyagamma` runs
+  −0.4–0.8% in the mean. augax uses the Gamma series.
+- **`4 < h < 8` with `|z| <= 4`**: `polyagamma` uses the saddlepoint
+  approximation, 1–2% low in the variance. augax uses the exact alternate
+  sampler.
+
+The saddlepoint and alternate samplers are ported from `polyagamma`
+(BSD-3-Clause). Every regime is tested against the closed-form mean and
+variance.
+
+The RNG is a per-call `xoshiro256**` seeded (via `splitmix64`) from the JAX PRNG
+key, so draws are reproducible and each call owns its stream — safe when several
+threads call at once. Normals and exponentials come from ziggurat samplers.
+
+An array longer than 65,536 elements is split into blocks of that size, each
+with its own stream, and the blocks are drawn in parallel on every hardware
+thread (set `AUGAX_NUM_THREADS` to cap it). The blocks depend only on the array
+length, so a draw is the same on any machine and any thread count. On a large
+Negative-Binomial model (8.1M cells) this takes the draw from 3.6 s to 0.33 s on
+16 threads.
+
+## The Kolmogorov scale
+
+A standard logistic disturbance is a normal scale mixture,
+`eps | lam ~ N(0, lam^2)` with `lam = 2K` and `K` Kolmogorov distributed
+(Stefanski 1991). `ks_sample(u, key)` draws `lam` given `eps = u`, which makes a
+logistic latent variable conditionally Gaussian: the propagated disturbance of a
+latent SAR logit, for one. The draw is exact: `1/lam^2` is proposed from an
+inverse Gaussian, `IG(1/|u|, 1)` (a chi-square at `|u| < 1e-4`), and accepted
+with probability `S(v) = f(v) e^{v/2} <= 1`, `f` the density of `(2K)^2`,
+evaluated lazily by the squeeze of an alternating series — the Kolmogorov
+series above `v = 4/3` and its Jacobi transform below, where the first's
+partial sums are not bounds (Holmes & Held 2006). It is checked against the
+identity `E[exp(-t^2 lam^2 / 2)] = pi t / sinh(pi t)` and against the
+posterior moments by quadrature. `ks_sample_numpy(u, rng)` is the same C++
+sampler for host callers.
+
+## Status
+
+- ✅ All `h > 0`: the logit (`h = 1`), Negative-Binomial (`h = y + alpha`) and
+  large-shape regimes.
+- Dispatch is automatic: the sampler is selected from `(h, z)` with no API
+  change.
+- ✅ `ks_sample`: exact for every `u`, on the device and on the host.
+
+CPU-only (the sampler is cheap scalar work; the point is to keep it on-device
+next to the rest of a JIT'd Gibbs sweep).