{"record":{"id":"b78d8a291398cbc9","repo":"jax-ml/jax","slug":"poisson-with-method-exact-is-only-implemented","errorCode":null,"errorMessage":"`poisson` with method='exact' is only implemented for the threefry2x32 RNG, not {key_impl}","messagePattern":"`poisson` with method='exact' is only implemented for the threefry2x32 RNG, not (.+?)","errorType":"validation","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/random/core.py","lineNumber":2223,"sourceCode":"  if shape is not None:\n    shape = core.canonicalize_shape(shape)\n  else:\n    shape = np.shape(lam)\n  out_sharding = canonicalize_sharding_for_samplers(out_sharding, \"poisson\", shape)\n  if method == 'approximate':\n    # don't preemptively broadcast lam, if lower rank it may save some computation\n    if lax.broadcast_shapes(np.shape(lam), shape) != shape:\n      raise ValueError(\"lam shape must be broadcastable to shape argument; \"\n                       f\"got lam.shape {np.shape(lam)}, shape {shape}\")\n    return maybe_auto_axes(_poisson_approx, out_sharding,\n                           shape=shape, dtype=dtype)(key, lam)\n  lam = jnp.broadcast_to(lam, shape)\n  # TODO(frostig): generalize underlying poisson implementation and\n  # remove this check\n  keys_dtype = typing.cast(prng.KeyTy, key.dtype)\n  key_impl = keys_dtype._impl\n  if key_impl is not threefry2x32.threefry_prng_impl:\n    raise NotImplementedError(\n        \"`poisson` with method='exact' is only implemented for the \"\n        f'threefry2x32 RNG, not {key_impl}')\n  lam = lax.convert_element_type(lam, np.float32)\n  return maybe_auto_axes(_poisson, out_sharding, shape=shape, dtype=dtype)(key, lam)\n\n\ndef gumbel(key: ArrayLike,\n           shape: Shape = (),\n           dtype: DTypeLikeFloat | None = None,\n           mode: str | None = None,\n           *,\n           out_sharding: NamedSharding | P | None = None) -> Array:\n  \"\"\"Sample Gumbel random values with given shape and float dtype.\n\n  The values are distributed according to the probability density function:\n\n  .. math::\n     f(x) = e^{-(x + e^{-x})}","sourceCodeStart":2205,"sourceCodeEnd":2241,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/random/core.py#L2205-L2241","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Use method='approximate', which works with any PRNG implementation.","Or create the key with the threefry implementation: jax.random.PRNGKey(seed) under default CPU/GPU config, or pass impl='threefry2x32' where supported.","Set the jax_default_prng_impl configuration back to threefry-compatible defaults if exactness matters more than TPU throughput."],"exampleFix":"// before\np = jax.random.poisson(key, lam)  # key is RBG -> NotImplementedError\n\n// after\np = jax.random.poisson(key, lam, method='approximate')","handlingStrategy":"fallback","validationCode":"import jax\nfrom jax._src import prng\nkey_impl = key.dtype._impl\nsupports_exact = key_impl is prng.threefry_prng_impl if hasattr(prng, 'threefry_prng_impl') else False\nmethod = 'exact' if supports_exact else 'approximate'","typeGuard":"def supports_exact_poisson(key) -> bool:\n    impl = key.dtype._impl\n    return 'threefry' in str(impl)","tryCatchPattern":"try:\n    p = jax.random.poisson(key, lam)  # method='exact'\nexcept NotImplementedError:\n    p = jax.random.poisson(key, lam, method='approximate')","preventionTips":["Check jax_default_prng_impl on TPU/RBG setups before using exact poisson.","Centralize key creation so the PRNG impl is known."],"tags":["jax","random","poisson","prng","not-implemented","tpu"],"backgroundTag":"unsupported-prng-implementation","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}