pypomp.random.fast_multinomial

pypomp.random.fast_multinomial(key: Array, n: Array, p: Array, order: int = 2, exact_max: int = 5, dtype: dtype | None = None) → Array[source]

Sample multinomial random variates using a GPU-optimized inverse CDF algorithm.

Generates multinomial counts by sequentially sampling binomial components via fast_binomial(). Follows the methodology from Giles and Beentjes (2024) [1]. Results are very close to exact but not guaranteed to be identical to a reference sampler.

Parameters:

  • key (jax.Array) – JAX PRNG key.

  • n (jax.Array) – Number of trials. Shape (...,).

  • p (jax.Array) – Category probabilities. Shape (..., k) where k is the number of categories. Probabilities along the last axis are normalised automatically.

  • order (int, optional) – Order of the beta-function approximation (0, 1, or 2). Defaults to 2 (most accurate).

  • exact_max (int, optional) – Maximum iterations for the bottom-up exact inverse CDF stage. Defaults to 5.

  • dtype (np.dtype or None, optional) – Output dtype (float or integer). Defaults to float64 if jax_enable_x64=True, otherwise float32. Integer dtypes return -1 for invalid inputs.

Returns:

Multinomial count array with the same shape as p and the specified dtype.

Return type:

jax.Array

Notes

For speed and accuracy metrics, see the Quant Tests.

References