API reference¶

The public API of pgjax is intentionally minimal — a single sampler function that works inside jax.jit, jax.vmap, lax.scan, and jax.pmap.

Sampling¶

pg_sample

Draw omega ~ PG(h, z) element-wise, on-device.

Internal helpers¶

These functions support pg_sample() and are documented for completeness; they are not part of the stable public API.

`_seed_from_key`	Derive an int64 RNG seed from a JAX PRNG key (per-key, reproducible).
`_dispatch`
`_sample_batched`
`_require_x64`

Module contents¶

pgjax: on-device Pólya-Gamma sampling as an XLA FFI custom call.

Draw omega ~ PG(h, z) from inside @jax.jit / lax.scan with no host round-trip — the native replacement for jax.pure_callback(random_polyagamma) that taxes every Gibbs sweep (and serializes under jax.pmap).

Exact for the Bernoulli/logit augmentation (h = 1) via the Devroye method (integer h = sum of Devroye draws), and for real-valued h — the Negative-Binomial h = y + alpha — via the tail-corrected Gamma-sum representation (validated bias-free vs polyagamma; see the README).

Example:

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

@jax.jit
def step(h, z, key):
    return pgjax.pg_sample(h, z, key)   # on-device PG draw

pgjax.pg_sample(h, z, key)[source]

Draw omega ~ PG(h, z) element-wise, on-device.

Parameters:

h (array_like) – Shape [n] float64. PG shape (h = 1 for Bernoulli/logit; positive integer for the exact sum-of-Devroye path).
z (array_like) – Shape [n] float64. PG tilt (the linear predictor).
key (jax.Array) – A JAX PRNG key. vmap over a batch of keys draws each element’s batch with its own RNG stream.

Returns:

omega – Draws from PG(h, z) with the same shape as h.

Return type:

jax.Array