jax-ml/jax · error · NotImplementedError

`poisson` with method='exact' is only implemented for the th

Error message

`poisson` with method='exact' is only implemented for the threefry2x32 RNG, not {key_impl}

What it means

The exact Poisson algorithm in JAX is implemented on top of the bit-level threefry2x32 PRNG and has not been generalized to other PRNG implementations. If your key uses a different impl (e.g. the default RBG/unsafe RBG keys on TPU, or a custom PRNG), jax.random.poisson with method='exact' (the default) raises NotImplementedError.

Source

Thrown at jax/_src/random/core.py:2223

  if shape is not None:
    shape = core.canonicalize_shape(shape)
  else:
    shape = np.shape(lam)
  out_sharding = canonicalize_sharding_for_samplers(out_sharding, "poisson", shape)
  if method == 'approximate':
    # don't preemptively broadcast lam, if lower rank it may save some computation
    if lax.broadcast_shapes(np.shape(lam), shape) != shape:
      raise ValueError("lam shape must be broadcastable to shape argument; "
                       f"got lam.shape {np.shape(lam)}, shape {shape}")
    return maybe_auto_axes(_poisson_approx, out_sharding,
                           shape=shape, dtype=dtype)(key, lam)
  lam = jnp.broadcast_to(lam, shape)
  # TODO(frostig): generalize underlying poisson implementation and
  # remove this check
  keys_dtype = typing.cast(prng.KeyTy, key.dtype)
  key_impl = keys_dtype._impl
  if key_impl is not threefry2x32.threefry_prng_impl:
    raise NotImplementedError(
        "`poisson` with method='exact' is only implemented for the "
        f'threefry2x32 RNG, not {key_impl}')
  lam = lax.convert_element_type(lam, np.float32)
  return maybe_auto_axes(_poisson, out_sharding, shape=shape, dtype=dtype)(key, lam)


def gumbel(key: ArrayLike,
           shape: Shape = (),
           dtype: DTypeLikeFloat | None = None,
           mode: str | None = None,
           *,
           out_sharding: NamedSharding | P | None = None) -> Array:
  """Sample Gumbel random values with given shape and float dtype.

  The values are distributed according to the probability density function:

  .. math::
     f(x) = e^{-(x + e^{-x})}

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use method='approximate', which works with any PRNG implementation.
2. Or create the key with the threefry implementation: jax.random.PRNGKey(seed) under default CPU/GPU config, or pass impl='threefry2x32' where supported.
3. Set the jax_default_prng_impl configuration back to threefry-compatible defaults if exactness matters more than TPU throughput.

Example fix

// before
p = jax.random.poisson(key, lam)  # key is RBG -> NotImplementedError

// after
p = jax.random.poisson(key, lam, method='approximate')

import jax
from jax._src import prng
key_impl = key.dtype._impl
supports_exact = key_impl is prng.threefry_prng_impl if hasattr(prng, 'threefry_prng_impl') else False
method = 'exact' if supports_exact else 'approximate'

def supports_exact_poisson(key) -> bool:
    impl = key.dtype._impl
    return 'threefry' in str(impl)

try:
    p = jax.random.poisson(key, lam)  # method='exact'
except NotImplementedError:
    p = jax.random.poisson(key, lam, method='approximate')

When it happens

Trigger: Calling jax.random.poisson(key, lam) (default method='exact') with a key from jax.random.PRNGKey on setups where the default PRNG is not threefry2x32 (e.g. jax_threefry_partitionable/TPU RBG defaults, or jax.make_key with an explicit impl), or any custom PRNG key.

Common situations: Running on TPU or with jax_default_prng_impl flag set to 'rbg'/'unsafe_rbg'; upgrading JAX versions where the default PRNG changed; custom sharded PRNG setups.

Related errors

• Cannot split a Pallas key. Use fold_in instead to generate n
• PRNG keys must be loaded from SMEM. Did you set the memory s
• Bit width must be 32
• QDWH implementation is only supported on TPU
• Failed to find assignment for logical_axis_index {logical_ax

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/b78d8a291398cbc9. Report an issue: GitHub.