API reference¶

The public API of augax is one sampler per augmentation scheme; each works inside jax.jit, jax.vmap, lax.scan, and jax.pmap.

Sampling¶

`pg_sample`	Draw `omega ~ PG(h, z)` element-wise, on-device.
`ks_sample`	Draw the Kolmogorov scale `lam` given a logistic disturbance `u`, on-device.
`ks_sample_numpy`	`ks_sample()` on the host: the same sampler, for NumPy callers.

Internal helpers¶

These functions support pg_sample() and ks_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`	`h`, `z` are `(B, ...)`, `seed` is `(B,)`.
`_ks_dispatch`
`_ks_sample_batched`	`u` is `(B, ...)`, `seed` is `(B,)`.
`_require_x64`

Module contents¶

augax: on-device draws for conditionally Gaussian augmentation, in JAX.

Each primitive is an XLA FFI custom call that samples one augmentation variable inside @jax.jit / lax.scan. It ships one:

pg_sample: Pólya–Gamma, omega ~ PG(h, z), for the logit and negative-binomial likelihoods (Polson, Scott & Windle 2013).
ks_sample: The Kolmogorov scale of a logistic disturbance, lam | eps, which makes the logistic a normal scale mixture (Stefanski 1991; Holmes & Held 2006). ks_sample_numpy() is the same draw for host (NumPy) callers.

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

Draws follow the hybrid scheme of polyagamma’s random_polyagamma (method=None): exact Devroye draws for the Bernoulli/logit augmentation (h = 1), the exact alternate sampler for other real h below the saddlepoint region — e.g. the Negative-Binomial h = y + alpha — a saddlepoint rejection sampler for larger h, and a moment-matched normal above h = 50 (see the README).

Example:

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

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

augax.ks_sample(u, key)[source]

Draw the Kolmogorov scale lam given a logistic disturbance u, on-device.

A standard logistic \(\varepsilon\) is the normal scale mixture \(\varepsilon \mid \lambda \sim N(0, \lambda^2)\), \(\lambda = 2K\) with \(K\) Kolmogorov distributed. Given \(\varepsilon = u\) this draws \(\lambda\) from

\[p(\lambda \mid u) \propto \lambda^{-1}\,\varphi(u/\lambda)\,f_{2K}(\lambda),\]

exactly: a rejection sampler with an inverse-Gaussian proposal for \(1/\lambda^2\) and the alternating-series squeeze of the Kolmogorov density in both of its forms (Holmes & Held 2006).

Parameters:

u (array_like) – Shape [n] float64: the disturbances.
key (jax.Array) – A JAX PRNG key. vmap over a batch of keys draws each element’s batch with its own RNG stream.

Returns:

lam – The scales, with the shape of u; lam**2 is the variance.

Return type:

jax.Array

augax.ks_sample_numpy(u, rng)[source]

ks_sample() on the host: the same sampler, for NumPy callers.

Parameters:

u (array_like) – The disturbances, any shape, as float64.
rng (numpy.random.Generator) – Supplies the seed of the C++ stream.

Returns:

lam – The scales, with the shape of u.

Return type:

numpy.ndarray

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