{"record":{"id":"d3d1c9780fbb3765","repo":"jax-ml/jax","slug":"alpha-must-be-one-dimensional-got-alpha-shape","errorCode":null,"errorMessage":"`alpha` must be one-dimensional; got alpha.shape={alpha.shape}","messagePattern":"`alpha` must be one-dimensional; got alpha\\.shape=(.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/scipy/stats/dirichlet.py","lineNumber":58,"sourceCode":"\n  where :math:`B(\\mathbf{\\alpha})` is the :func:`~jax.scipy.special.beta` function\n  in a :math:`K`-dimensional vector space.\n\n  Args:\n    x: arraylike, value at which to evaluate the PDF\n    alpha: arraylike, distribution shape parameter\n\n  Returns:\n    array of logpdf values.\n\n  See Also:\n    :func:`jax.scipy.stats.dirichlet.pdf`\n  \"\"\"\n  return _logpdf(*promote_dtypes_inexact(x, alpha))\n\ndef _logpdf(x: Array, alpha: Array) -> Array:\n  if alpha.ndim != 1:\n    raise ValueError(\n      f\"`alpha` must be one-dimensional; got alpha.shape={alpha.shape}\"\n    )\n  if x.shape[0] not in (alpha.shape[0], alpha.shape[0] - 1):\n    raise ValueError(\n      \"`x` must have either the same number of entries as `alpha` \"\n      f\"or one entry fewer; got x.shape={x.shape}, alpha.shape={alpha.shape}\"\n    )\n  one = _lax_const(x, 1)\n  if x.shape[0] != alpha.shape[0]:\n    x = jnp.concatenate([x, lax.sub(one, x.sum(0, keepdims=True))], axis=0)\n  normalize_term = jnp.sum(gammaln(alpha)) - gammaln(jnp.sum(alpha))\n  if x.ndim > 1:\n    alpha = lax.broadcast_in_dim(alpha, alpha.shape + (1,) * (x.ndim - 1), (0,))\n  log_probs = lax.sub(jnp.sum(xlogy(lax.sub(alpha, one), x), axis=0), normalize_term)\n  return jnp.where(_is_simplex(x), log_probs, -np.inf)\n\n\ndef pdf(x: ArrayLike, alpha: ArrayLike) -> Array:","sourceCodeStart":40,"sourceCodeEnd":76,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/scipy/stats/dirichlet.py#L40-L76","documentation":"The Dirichlet logpdf/pdf in jax.scipy.stats.dirichlet requires the concentration parameter alpha to be a 1-D array. Because the check is on alpha.ndim, passing a batched 2-D alpha (e.g. shape (batch, k)) raises immediately.","triggerScenarios":"Calling jax.scipy.stats.dirichlet.logpdf(x, alpha) with alpha of shape (2, 3) or any ndim != 1, e.g. when batched alphas were kept as a matrix.","commonSituations":"Vectorizing over multiple Dirichlet distributions and passing a stacked alpha matrix instead of using vmap; porting code where scipy tolerated broadcasting (scipy also errors, but users assume batch support).","solutions":["Squeeze/reshape alpha to 1-D before the call: alpha = alpha.squeeze()","Use jax.vmap(jax.scipy.stats.dirichlet.logpdf, in_axes=(0, 0)) to batch over distributions","Verify you did not accidentally pass x and alpha in swapped order"],"exampleFix":"// before\nlp = jax.scipy.stats.dirichlet.logpdf(x, alpha)  # alpha.shape == (B, K)\n// after\nlp = jax.vmap(jax.scipy.stats.dirichlet.logpdf)(x, alpha)  # per-sample alpha of shape (K,)","handlingStrategy":"validation","validationCode":"import jax.numpy as jnp\nassert jnp.asarray(alpha).ndim == 1, 'alpha must be 1-D'","typeGuard":"def is_1d_alpha(alpha) -> bool:\n    return jnp.asarray(alpha).ndim == 1","tryCatchPattern":null,"preventionTips":["Normalize alpha to a flat vector at call sites","Use vmap for batched Dirichlet parameters","Add shape asserts in test fixtures"],"tags":["jax","scipy","dirichlet","shape-validation"],"backgroundTag":"invalid-array-shape","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}