jax-ml/jax · error · ValueError

`alpha` must be one-dimensional; got alpha.shape={alpha.shap

Error message

`alpha` must be one-dimensional; got alpha.shape={alpha.shape}

What it means

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.

Source

Thrown at jax/_src/scipy/stats/dirichlet.py:58

  where :math:`B(\mathbf{\alpha})` is the :func:`~jax.scipy.special.beta` function
  in a :math:`K`-dimensional vector space.

  Args:
    x: arraylike, value at which to evaluate the PDF
    alpha: arraylike, distribution shape parameter

  Returns:
    array of logpdf values.

  See Also:
    :func:`jax.scipy.stats.dirichlet.pdf`
  """
  return _logpdf(*promote_dtypes_inexact(x, alpha))

def _logpdf(x: Array, alpha: Array) -> Array:
  if alpha.ndim != 1:
    raise ValueError(
      f"`alpha` must be one-dimensional; got alpha.shape={alpha.shape}"
    )
  if x.shape[0] not in (alpha.shape[0], alpha.shape[0] - 1):
    raise ValueError(
      "`x` must have either the same number of entries as `alpha` "
      f"or one entry fewer; got x.shape={x.shape}, alpha.shape={alpha.shape}"
    )
  one = _lax_const(x, 1)
  if x.shape[0] != alpha.shape[0]:
    x = jnp.concatenate([x, lax.sub(one, x.sum(0, keepdims=True))], axis=0)
  normalize_term = jnp.sum(gammaln(alpha)) - gammaln(jnp.sum(alpha))
  if x.ndim > 1:
    alpha = lax.broadcast_in_dim(alpha, alpha.shape + (1,) * (x.ndim - 1), (0,))
  log_probs = lax.sub(jnp.sum(xlogy(lax.sub(alpha, one), x), axis=0), normalize_term)
  return jnp.where(_is_simplex(x), log_probs, -np.inf)


def pdf(x: ArrayLike, alpha: ArrayLike) -> Array:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Squeeze/reshape alpha to 1-D before the call: alpha = alpha.squeeze()
2. Use jax.vmap(jax.scipy.stats.dirichlet.logpdf, in_axes=(0, 0)) to batch over distributions
3. Verify you did not accidentally pass x and alpha in swapped order

Example fix

// before
lp = jax.scipy.stats.dirichlet.logpdf(x, alpha)  # alpha.shape == (B, K)
// after
lp = jax.vmap(jax.scipy.stats.dirichlet.logpdf)(x, alpha)  # per-sample alpha of shape (K,)

import jax.numpy as jnp
assert jnp.asarray(alpha).ndim == 1, 'alpha must be 1-D'

def is_1d_alpha(alpha) -> bool:
    return jnp.asarray(alpha).ndim == 1

When it happens

Trigger: 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.

Common situations: 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).

Related errors

• `x` must have either the same number of entries as `alpha` o
• tensorinv is only possible when the product of the first `in
• zero-size arrays not supported in convolutions, got shapes {
• convolve2d() only supports 2-dimensional inputs.
• correlate2d() only supports 2-dimensional inputs.

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