jax-ml/jax · error · ValueError

`x` must have either the same number of entries as `alpha` o

Error message

`x` must have either the same number of entries as `alpha` or one entry fewer; got x.shape={x.shape}, alpha.shape={alpha.shape}

What it means

For jax.scipy.stats.dirichlet.logpdf/pdf, x along axis 0 must have exactly len(alpha) entries (full simplex) or len(alpha)-1 entries (last coordinate implicit, computed as 1 - sum(x)). Any other leading dimension raises this error.

Source

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

  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:
  r"""Dirichlet probability distribution function.

  JAX implementation of :obj:`scipy.stats.dirichlet` ``pdf``.

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Reshape x so x.shape[0] is alpha.shape[0] or alpha.shape[0]-1
2. If the last simplex coordinate was dropped, keep only one dropped value; supply the full k-length x if you dropped more
3. Put batch dimensions on axes other than axis 0 or use vmap

Example fix

// before
alpha = jnp.array([2.0, 3.0, 4.0])
x = jnp.array([0.5])  # wrong length
lp = jax.scipy.stats.dirichlet.logpdf(x, alpha)
// after
x = jnp.array([0.5, 0.3])  # k-1 entries; last = 1 - 0.8
lp = jax.scipy.stats.dirichlet.logpdf(x, alpha)

k = alpha.shape[0]
assert x.shape[0] in (k, k - 1), f'x.shape[0] must be {k} or {k-1}'

def dirichlet_x_valid(x, alpha) -> bool:
    return jnp.asarray(x).shape[0] in (jnp.asarray(alpha).shape[0], jnp.asarray(alpha).shape[0] - 1)

When it happens

Trigger: Passing x with shape (k+2, ...) or (k-3, ...) when alpha has k entries; passing x whose batch axis is on axis 0 instead of matching alpha's length.

Common situations: Feeding unbatched x of wrong length; confusing the batch dimension with the category dimension; forgetting that the implicit form drops exactly one coordinate, not an arbitrary number.

Related errors

• `alpha` must be one-dimensional; got alpha.shape={alpha.shap
• 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/1af2dc0d71b02b36. Report an issue: GitHub.