jax-ml/jax · error · ValueError

lam shape must be broadcastable to shape argument; got lam.s

Error message

lam shape must be broadcastable to shape argument; got lam.shape {np.shape(lam)}, shape {shape}

What it means

When jax.random.poisson is called with method='approximate' and an explicit shape, JAX does not preemptively broadcast lam; instead it verifies that lam broadcasts up to exactly shape. If broadcasting lam to shape would change the result shape (i.e. lax.broadcast_shapes differs), it raises ValueError naming both shapes.

Source

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

  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
  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 = (),

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Make lam broadcast-compatible with shape: lam of shape (), (1,), or exactly the trailing dims of shape, e.g. poisson(key, lam[None, :], shape=(4, 3)).
2. Or broadcast lam explicitly first: lam = jnp.broadcast_to(lam, shape) then omit shape / pass matching shape.
3. Check np.shape(lam) vs the intended shape with a quick assert before calling in shape-sensitive pipelines.

Example fix

// before
p = jax.random.poisson(key, lam, shape=(4, 3), method='approximate')  # lam.shape == (5,)

// after
p = jax.random.poisson(key, lam[:3], shape=(4, 3), method='approximate')  # broadcastable

import jax.numpy as jnp, numpy as np
if method == 'approximate' and shape is not None:
    assert jnp.broadcast_shapes(np.shape(lam), tuple(shape)) == tuple(shape)

def lam_broadcasts_to(lam, shape) -> bool:
    import numpy as np
    try:
        return np.broadcast_shapes(np.shape(lam), tuple(shape)) == tuple(shape)
    except ValueError:
        return False

When it happens

Trigger: jax.random.poisson(key, lam, shape=(4, 3)) with lam.shape == (5,) (incompatible), or lam.shape == (3, 5) with shape == (4, 3) where broadcasting produces something other than shape; only with method='approximate' and shape is not None.

Common situations: Assuming the sampler silently truncates or reshapes lam to shape; passing batch shape that does not account for an extra leading dimension in lam; refactoring shapes upstream so lam gained/lost a batch axis.

Related errors

• dirichlet requires alpha.ndim >= 1, got alpha.ndim == {}
• method argument to `poisson` must be one of {'exact', 'appro
• `poisson` with method='exact' is only implemented for the th
• Number of samples without replacement ({k}) cannot exceed nu
• {name} was requested to map its argument along axis {axis},

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