{"record":{"id":"518008b47cd90804","repo":"jax-ml/jax","slug":"lam-shape-must-be-broadcastable-to-shape-argument","errorCode":null,"errorMessage":"lam shape must be broadcastable to shape argument; got lam.shape {np.shape(lam)}, shape {shape}","messagePattern":"lam shape must be broadcastable to shape argument; got lam\\.shape (.+?), shape (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/random/core.py","lineNumber":2213,"sourceCode":"  Returns:\n    A random array with the specified dtype and with shape given by ``shape`` if\n    ``shape is not None, or else by ``lam.shape``.\n  \"\"\"\n  key, _ = _check_prng_key(\"poisson\", key)\n  if method not in {'exact', 'approximate'}:\n    raise ValueError(\"method argument to `poisson` must be one of \"\n                     f\"{{'exact', 'approximate'}}, got {method!r}\")\n  dtype = dtypes.check_and_canonicalize_user_dtype(\n      int if dtype is None else dtype)\n  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 = (),","sourceCodeStart":2195,"sourceCodeEnd":2231,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/random/core.py#L2195-L2231","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["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)).","Or broadcast lam explicitly first: lam = jnp.broadcast_to(lam, shape) then omit shape / pass matching shape.","Check np.shape(lam) vs the intended shape with a quick assert before calling in shape-sensitive pipelines."],"exampleFix":"// before\np = jax.random.poisson(key, lam, shape=(4, 3), method='approximate')  # lam.shape == (5,)\n\n// after\np = jax.random.poisson(key, lam[:3], shape=(4, 3), method='approximate')  # broadcastable","handlingStrategy":"validation","validationCode":"import jax.numpy as jnp, numpy as np\nif method == 'approximate' and shape is not None:\n    assert jnp.broadcast_shapes(np.shape(lam), tuple(shape)) == tuple(shape)","typeGuard":"def lam_broadcasts_to(lam, shape) -> bool:\n    import numpy as np\n    try:\n        return np.broadcast_shapes(np.shape(lam), tuple(shape)) == tuple(shape)\n    except ValueError:\n        return False","tryCatchPattern":null,"preventionTips":["Prefer letting shape default to lam.shape when possible.","Broadcast lam explicitly with jnp.broadcast_to before calling."],"tags":["jax","random","poisson","broadcasting","shape"],"backgroundTag":"shape-broadcast-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}