jax-ml/jax · error · ValueError

method argument to `poisson` must be one of {'exact', 'appro

Error message

method argument to `poisson` must be one of {'exact', 'approximate'}, got {method!r}

What it means

jax.random.poisson accepts only method='exact' (algorithm restricted to the threefry2x32 PRNG, see the separate NotImplementedError) or method='approximate' (a normal/knuth-style approximation that works on any backend). Any other method string raises ValueError.

Source

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

      sampling and supports only the threefry2x32 RNG. The ``'approximate'``
      method is loop-free and faster but approximate: the total variation
      distance from the exact distribution is below 1e-4.
    out_sharding: Optional. Specifies how the output array should be sharded
      across devices in multi-device computation. Can be a
      :class:`~jax.sharding.NamedSharding`, a :class:`~jax.sharding.PartitionSpec`
      (``P``), or ``None`` (default). When specified, the output will be sharded
      according to the given sharding specification. Primarily used in explicit
      sharding mode.
      See the `explicit sharding tutorial <https://docs.jax.dev/en/latest/parallel.html>`_
      for more details.

  Returns:
    A random array with the specified dtype and with shape given by ``shape`` if
    ``shape is not None, or else by ``lam.shape``.
  """
  key, _ = _check_prng_key("poisson", key)
  if method not in {'exact', 'approximate'}:
    raise ValueError("method argument to `poisson` must be one of "
                     f"{{'exact', 'approximate'}}, got {method!r}")
  dtype = dtypes.check_and_canonicalize_user_dtype(
      int if dtype is None else dtype)
  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

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Use exactly 'exact' or 'approximate'.
If you hit this while migrating off the default, note method='approximate' avoids the threefry2x32-only limitation of 'exact'.
Validate method strings from config at load time.

Example fix

// before
p = jax.random.poisson(key, 3.0, method='normal')

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

Defensive patterns

Strategy: validation

Validate before calling

method = method or 'exact'
assert method in {'exact', 'approximate'}, 'poisson method must be exact or approximate'

Type guard

def is_valid_poisson_method(method: str) -> bool:
    return method in {'exact', 'approximate'}

Prevention

On non-threefry PRNGs, plan for method='approximate'.
Validate method from config against the allowed set.

When it happens

Trigger: Calling jax.random.poisson(key, lam, method='Normal') or any string not in {'exact','approximate'}; explicitly passing method=None.

Common situations: Config-driven method names shared across gamma/poisson; typos; users switching to method='approximate' when they hit the threefry2x32 restriction but misspelling it and getting this error instead.

Related errors

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