pgjax

On-device Pólya-Gamma sampling for JAX, as an XLA FFI custom call.

pgjax.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.

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

@jax.jit
def omega_update(h, z, key):
    return pgjax.pg_sample(h, z, key)

Method

  • Small integer h (incl. Bernoulli/logit h = 1): the exact Devroye sampler for PG(1, z) (Polson, Scott & Windle 2013), summed h times.
  • Real h (e.g. Negative-Binomial h = y + alpha): the exact infinite-sum (Gamma) representation truncated at K terms with a moment-matched Gamma tail correction for the remainder. The tail mean/variance have closed forms (Σ 1/d_k = (π/2a)tanh(πa), and its a-derivative), so the correction cancels the truncation bias in the mean — it is precisely that bias (truncation with no tail term) that collapses alpha in NB models when omitted.
  • Large h: the saddlepoint rejection sampler (Windle, Polson & Scott 2014, arXiv:1405.0506). A two-piece bounding-kernel envelope proposes, and the saddlepoint density approximation accepts, giving O(1) per accepted draw versus the O(K) Gamma sum. Used automatically when h >= 8 or (h > 4 and |z| <= 4) — the hybrid dispatch thresholds validated in zoj613/polyagamma.

Validated against random_polyagamma over h ∈ [0.5, 40] × z ∈ [-8, 8]: KS ≤ 0.009, relative mean error ≤ 0.34%. The saddlepoint path is validated to KS < 0.02 and relative mean error < 2% over h ∈ {8,…,40} × z ∈ {0,1,-4,5}.

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 under the concurrent per-device calls jax.pmap issues.

Status

  • ✅ Exact PG(1, z) / integer h (Devroye), real h (tail-corrected Gamma sum), and large h (saddlepoint) — covers the logit (h=1), Negative-Binomial (h=y+alpha), and large-shape regimes.
  • Hybrid dispatch is automatic: the appropriate sampler is selected from (h, z) with no API change.

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).

Key dates

PyPI data
Data sourced directly from PyPI's database.
  • Released:
Latest release

1 maintainer

PyPI data
Data sourced directly from PyPI's database.
Avatar for knaaptime from gravatar.com knaaptime

License

BSD-3

Requires

Python >=3.9

Provides Extra

test

Report project as malware

Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

pgjax-0.1.0.tar.gz (22.7 kB view details)

Uploaded Source

Built Distributions

If you're not sure about the file name format, learn more about wheel file names.

pgjax-0.1.0-cp313-cp313-win_amd64.whl (245.6 kB view details)

Uploaded CPython 3.13Windows x86-64

pgjax-0.1.0-cp313-cp313-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl (59.7 kB view details)

Uploaded CPython 3.13manylinux: glibc 2.24+ x86-64manylinux: glibc 2.28+ x86-64

pgjax-0.1.0-cp313-cp313-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl (55.9 kB view details)

Uploaded CPython 3.13manylinux: glibc 2.24+ ARM64manylinux: glibc 2.28+ ARM64

pgjax-0.1.0-cp313-cp313-macosx_15_0_arm64.whl (51.3 kB view details)

Uploaded CPython 3.13macOS 15.0+ ARM64

pgjax-0.1.0-cp312-cp312-win_amd64.whl (245.7 kB view details)

Uploaded CPython 3.12Windows x86-64

pgjax-0.1.0-cp312-cp312-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl (59.7 kB view details)

Uploaded CPython 3.12manylinux: glibc 2.24+ x86-64manylinux: glibc 2.28+ x86-64

pgjax-0.1.0-cp312-cp312-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl (55.9 kB view details)

Uploaded CPython 3.12manylinux: glibc 2.24+ ARM64manylinux: glibc 2.28+ ARM64

pgjax-0.1.0-cp312-cp312-macosx_15_0_arm64.whl (51.3 kB view details)

Uploaded CPython 3.12macOS 15.0+ ARM64

pgjax-0.1.0-cp311-cp311-win_amd64.whl (246.7 kB view details)

Uploaded CPython 3.11Windows x86-64

pgjax-0.1.0-cp311-cp311-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl (60.8 kB view details)

Uploaded CPython 3.11manylinux: glibc 2.24+ x86-64manylinux: glibc 2.28+ x86-64

pgjax-0.1.0-cp311-cp311-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl (57.3 kB view details)

Uploaded CPython 3.11manylinux: glibc 2.24+ ARM64manylinux: glibc 2.28+ ARM64

pgjax-0.1.0-cp311-cp311-macosx_15_0_arm64.whl (52.5 kB view details)

Uploaded CPython 3.11macOS 15.0+ ARM64

File details

Details for the file pgjax-0.1.0.tar.gz.

File metadata

  • Download URL: pgjax-0.1.0.tar.gz
  • Upload date:
  • Size: 22.7 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? Yes
  • Uploaded via: twine/7.0.0 CPython/3.13.14

File hashes

Hashes for pgjax-0.1.0.tar.gz
Algorithm Hash digest
SHA256 3cb526aeab1c26a2c588e327ea4115a01d337f23f828d14aa3e4b97581704408
MD5 240695ebb3290aba271d9973428d0485
BLAKE2b-256 cc7f8a4bffd3656f2c99dad46d6ba8ec4f4f8aa5584d02542e1ca4d8314088b3

See more details on using hashes here.

Provenance

The following attestation bundles were made for pgjax-0.1.0.tar.gz:

Publisher: build-wheels.yml on knaaptime/pgjax

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

File details

Details for the file pgjax-0.1.0-cp313-cp313-win_amd64.whl.

File metadata

  • Download URL: pgjax-0.1.0-cp313-cp313-win_amd64.whl
  • Upload date:
  • Size: 245.6 kB
  • Tags: CPython 3.13, Windows x86-64
  • Uploaded using Trusted Publishing? Yes
  • Uploaded via: twine/7.0.0 CPython/3.13.14

File hashes

Hashes for pgjax-0.1.0-cp313-cp313-win_amd64.whl
Algorithm Hash digest
SHA256 10eeadc1bc79b86623b9b99dbf529d6d7fbac3cafbfdb6b22113dddc62aec93a
MD5 4f024f2af18ba6a55c72c0e7858d3681
BLAKE2b-256 935984491b64d32a7b8f83433048eaec52c6454a3c1f6e0193c0acd19abfa0e1

See more details on using hashes here.

Provenance

The following attestation bundles were made for pgjax-0.1.0-cp313-cp313-win_amd64.whl:

Publisher: build-wheels.yml on knaaptime/pgjax

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

File details

Details for the file pgjax-0.1.0-cp313-cp313-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl.

File metadata

File hashes

Hashes for pgjax-0.1.0-cp313-cp313-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl
Algorithm Hash digest
SHA256 1f5b96066d0661b268fcd19c001e0e2492572a3ccf4310541969d89fadb7db0a
MD5 4be182b5567dff2bddb724cd757653fa
BLAKE2b-256 d235015bf4679c6d9c5648b26250934985d0fc59606d18efb250ade203888444

See more details on using hashes here.

Provenance

The following attestation bundles were made for pgjax-0.1.0-cp313-cp313-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl:

Publisher: build-wheels.yml on knaaptime/pgjax

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

File details

Details for the file pgjax-0.1.0-cp313-cp313-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl.

File metadata

File hashes

Hashes for pgjax-0.1.0-cp313-cp313-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 629a6c902b0c5a470211058b6d9817a6cf3c9b4e6a4498a72b2377eb2641874b
MD5 bdef9b96c3fdb73601ab903a67dd8cfb
BLAKE2b-256 760e454b1e8c05a2914add8862bb4b212cffaffef3a394cd9cd834ad0b7ca0f8

See more details on using hashes here.

Provenance

The following attestation bundles were made for pgjax-0.1.0-cp313-cp313-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl:

Publisher: build-wheels.yml on knaaptime/pgjax

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

File details

Details for the file pgjax-0.1.0-cp313-cp313-macosx_15_0_arm64.whl.

File metadata

File hashes

Hashes for pgjax-0.1.0-cp313-cp313-macosx_15_0_arm64.whl
Algorithm Hash digest
SHA256 93b3a04316e5d9f83af0a66178d750ebd9df488bf701f8adce39cb54bcfb2f52
MD5 f53a52527de3348cb1f218949e670dcd
BLAKE2b-256 407cec45be5c739be05a4e22d2718b7f89da745e658e9cbd28dc655a0a3fe44e

See more details on using hashes here.

Provenance

The following attestation bundles were made for pgjax-0.1.0-cp313-cp313-macosx_15_0_arm64.whl:

Publisher: build-wheels.yml on knaaptime/pgjax

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

File details

Details for the file pgjax-0.1.0-cp312-cp312-win_amd64.whl.

File metadata

  • Download URL: pgjax-0.1.0-cp312-cp312-win_amd64.whl
  • Upload date:
  • Size: 245.7 kB
  • Tags: CPython 3.12, Windows x86-64
  • Uploaded using Trusted Publishing? Yes
  • Uploaded via: twine/7.0.0 CPython/3.13.14

File hashes

Hashes for pgjax-0.1.0-cp312-cp312-win_amd64.whl
Algorithm Hash digest
SHA256 ab4913f58f8e72026bd759c4f44afab40d8bffa3b67d49bf53b96e3972990782
MD5 cac22a1568182e297c96ed33559e2b1c
BLAKE2b-256 d64968d37b25ae40830b6acdc0db0b9da078c9047dd6df91c532e22fcc51aae9

See more details on using hashes here.

Provenance

The following attestation bundles were made for pgjax-0.1.0-cp312-cp312-win_amd64.whl:

Publisher: build-wheels.yml on knaaptime/pgjax

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

File details

Details for the file pgjax-0.1.0-cp312-cp312-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl.

File metadata

File hashes

Hashes for pgjax-0.1.0-cp312-cp312-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl
Algorithm Hash digest
SHA256 248c48214c8c1d7bed6807683e10a7264872556cd5b9d683601cdaf4161366b6
MD5 7217f13f15a22c00ddc6a144f38d4044
BLAKE2b-256 f3e701cd31f98bdb07145a5525839a516ed280b594b0a18930af789b58e09f12

See more details on using hashes here.

Provenance

The following attestation bundles were made for pgjax-0.1.0-cp312-cp312-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl:

Publisher: build-wheels.yml on knaaptime/pgjax

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

File details

Details for the file pgjax-0.1.0-cp312-cp312-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl.

File metadata

File hashes

Hashes for pgjax-0.1.0-cp312-cp312-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 b025bb10c1f58597e215db6db8b80bba23ce01f181f3255d21656c56dccfaaea
MD5 8db789cb9204973d1e11609ba61a8ab5
BLAKE2b-256 9e303322642e7dd4459b17c6fc8455b055cf52c38319f5265bc65027a1404389

See more details on using hashes here.

Provenance

The following attestation bundles were made for pgjax-0.1.0-cp312-cp312-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl:

Publisher: build-wheels.yml on knaaptime/pgjax

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

File details

Details for the file pgjax-0.1.0-cp312-cp312-macosx_15_0_arm64.whl.

File metadata

File hashes

Hashes for pgjax-0.1.0-cp312-cp312-macosx_15_0_arm64.whl
Algorithm Hash digest
SHA256 11ff8ab81ac5b6f50dfedbb65761d781da418d89fa0e220437bf7b00426e91d4
MD5 03f8c9651c8a3b3daa87251b07571117
BLAKE2b-256 8b299f616b4d82724d858586c6a9d9af99b4f0f2dc3fa027f017c7e5a7e8b770

See more details on using hashes here.

Provenance

The following attestation bundles were made for pgjax-0.1.0-cp312-cp312-macosx_15_0_arm64.whl:

Publisher: build-wheels.yml on knaaptime/pgjax

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

File details

Details for the file pgjax-0.1.0-cp311-cp311-win_amd64.whl.

File metadata

  • Download URL: pgjax-0.1.0-cp311-cp311-win_amd64.whl
  • Upload date:
  • Size: 246.7 kB
  • Tags: CPython 3.11, Windows x86-64
  • Uploaded using Trusted Publishing? Yes
  • Uploaded via: twine/7.0.0 CPython/3.13.14

File hashes

Hashes for pgjax-0.1.0-cp311-cp311-win_amd64.whl
Algorithm Hash digest
SHA256 b5dfade87b7e008f42c1619f2aa1fa6cd995af3581fd0496b9a0103d730f6925
MD5 ccafccc8830ca133d6090f6fa2a4f88b
BLAKE2b-256 943e7742547be910216b9c18485429a0267eba724d2eb1f8998b49a3e4dba7f9

See more details on using hashes here.

Provenance

The following attestation bundles were made for pgjax-0.1.0-cp311-cp311-win_amd64.whl:

Publisher: build-wheels.yml on knaaptime/pgjax

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

File details

Details for the file pgjax-0.1.0-cp311-cp311-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl.

File metadata

File hashes

Hashes for pgjax-0.1.0-cp311-cp311-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl
Algorithm Hash digest
SHA256 e5e46d1601d2ba6527a11f77122fa9a1fbb4b4daf4f7513000a5bfe9cdd9c48b
MD5 3b1d67ed76b8c40914eb24ecffa8a6f0
BLAKE2b-256 7df14d66bf4514a9c8a57d6927b6dcbfa4286e4b63ff51bc26264782ba79b857

See more details on using hashes here.

Provenance

The following attestation bundles were made for pgjax-0.1.0-cp311-cp311-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl:

Publisher: build-wheels.yml on knaaptime/pgjax

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

File details

Details for the file pgjax-0.1.0-cp311-cp311-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl.

File metadata

File hashes

Hashes for pgjax-0.1.0-cp311-cp311-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 a292f28bf2928f64b64f24c2f83c5c64ef76b4b36ebdecdcae4a3e666d9b5eeb
MD5 518c750f63e2be246816847ad89754b0
BLAKE2b-256 89bc012c73aa5a2b28032dc88fb707968b7aa5b6b17bf95cd2c9d82653a11189

See more details on using hashes here.

Provenance

The following attestation bundles were made for pgjax-0.1.0-cp311-cp311-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl:

Publisher: build-wheels.yml on knaaptime/pgjax

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

File details

Details for the file pgjax-0.1.0-cp311-cp311-macosx_15_0_arm64.whl.

File metadata

File hashes

Hashes for pgjax-0.1.0-cp311-cp311-macosx_15_0_arm64.whl
Algorithm Hash digest
SHA256 f8d50985097c7cd6765b9957d177a3ed4e06d3c3f249e26057db3eb01e4f2ed9
MD5 75adda3554224aa6f34501787846f945
BLAKE2b-256 6d05030da0550f080e46d0ce95ae6ad093b4378749a76cd00ca287fb407ccbb3

See more details on using hashes here.

Provenance

The following attestation bundles were made for pgjax-0.1.0-cp311-cp311-macosx_15_0_arm64.whl:

Publisher: build-wheels.yml on knaaptime/pgjax

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

Release history Release notifications | RSS feed

This release

0.1.0 This release

13 files

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Depot Continuous Integration Fastly CDN Google Download Analytics Sentry Error logging StatusPage Status page