{"record":{"id":"077aa49eb851936b","repo":"jax-ml/jax","slug":"dirichlet-requires-alpha-ndim-1-got-alpha-ndim","errorCode":null,"errorMessage":"dirichlet requires alpha.ndim >= 1, got alpha.ndim == {}","messagePattern":"dirichlet requires alpha\\.ndim >= 1, got alpha\\.ndim == (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/random/core.py","lineNumber":1422,"sourceCode":"  key, _ = _check_prng_key(\"dirichlet\", key)\n  dtype = dtypes.check_and_canonicalize_user_dtype(\n      float if dtype is None else dtype)\n  if not dtypes.issubdtype(dtype, np.floating):\n    raise ValueError(f\"dtype argument to `dirichlet` must be a float \"\n                     f\"dtype, got {dtype}\")\n  if shape is not None:\n    shape = core.canonicalize_shape(shape)\n  out_sharding = canonicalize_sharding_for_samplers(out_sharding, \"dirichlet\", shape)\n  return maybe_auto_axes(_dirichlet, out_sharding,\n                         shape=shape, dtype=dtype)(key, alpha)\n\n@jit(static_argnums=(2, 3))\ndef _dirichlet(key, alpha, shape, dtype) -> Array:\n  from jax._src.nn.functions import softmax  # pyrefly: ignore[missing-import]\n\n  if not np.ndim(alpha) >= 1:\n    msg = \"dirichlet requires alpha.ndim >= 1, got alpha.ndim == {}\"\n    raise ValueError(msg.format(np.ndim(alpha)))\n\n  if shape is None:\n    shape = np.shape(alpha)[:-1]\n  else:\n    _check_shape(\"dirichlet\", shape, np.shape(alpha)[:-1])\n\n  alpha = lax.convert_element_type(alpha, dtype)\n\n  # Compute gamma in log space, otherwise small alpha can lead to poor behavior.\n  log_gamma_samples = loggamma(key, alpha, shape + np.shape(alpha)[-1:], dtype)\n  return softmax(log_gamma_samples, -1)\n\n\ndef exponential(key: ArrayLike,\n                shape: Shape = (),\n                dtype: DTypeLikeFloat | None = None,\n                *,\n                out_sharding: NamedSharding | P | None = None) -> Array:","sourceCodeStart":1404,"sourceCodeEnd":1440,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/random/core.py#L1404-L1440","documentation":"jax.random.dirichlet requires the concentration parameter alpha to have at least one dimension, because the last axis of alpha holds the category dimension that the Dirichlet distribution normalizes over. A scalar alpha (ndim == 0) has no category axis, so the sampler cannot form a valid probability simplex. JAX therefore raises ValueError before tracing the jitted _dirichlet implementation.","triggerScenarios":"Calling jax.random.dirichlet(key, alpha) where alpha is a Python scalar, a 0-d jnp array, or np.ndim(alpha) == 0 (e.g. jax.random.dirichlet(key, 1.0) instead of jax.random.dirichlet(key, jnp.array([1.0]))).","commonSituations":"Porting NumPy/SciPy code where a scalar was accepted as a single-category concentration; building alpha from a computation that accidentally reduces to a scalar (e.g. taking [-1] indexing or a squeeze/mean); passing an unshaped parameter from a config dataclass.","solutions":["Give alpha a category axis: pass an array of shape (k,) or (..., k), e.g. jnp.ones((k,)) or jnp.atleast_1d(alpha).","If alpha arrives from user config, normalize it at the boundary with jnp.atleast_1d(jnp.asarray(alpha, dtype=float)).","Check np.ndim(alpha) >= 1 before calling dirichlet and raise a clearer domain-specific error."],"exampleFix":"// before\nsamples = jax.random.dirichlet(key, 2.0)  # scalar -> ValueError\n\n// after\nsamples = jax.random.dirichlet(key, jnp.full((5,), 2.0))  # shape (5,) concentrations","handlingStrategy":"validation","validationCode":"import numpy as np, jax.numpy as jnp\nalpha = jnp.atleast_1d(jnp.asarray(alpha, dtype=float))\nassert np.ndim(alpha) >= 1","typeGuard":"def is_valid_dirichlet_alpha(alpha) -> bool:\n    return np.ndim(alpha) >= 1","tryCatchPattern":null,"preventionTips":["Always construct alpha as a 1+-dim array, e.g. jnp.full((k,), val).","Wrap external config inputs with jnp.atleast_1d before sampling."],"tags":["jax","random","dirichlet","input-validation","shape"],"backgroundTag":"invalid-shape-argument","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}