pgjax.pg_sample¶

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