What causes "dirichlet requires alpha.ndim >= 1, got alpha.ndim == {}" in jax-ml/jax?

Give alpha a category axis: pass an array of shape (k,) or (..., k), e.g. jnp.ones((k,)) or jnp.atleast_1d(alpha).

jax-ml/jax · error · ValueError

dirichlet requires alpha.ndim >= 1, got alpha.ndim == {}

Error message

dirichlet requires alpha.ndim >= 1, got alpha.ndim == {}

What it means

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.

Source

Thrown at jax/_src/random/core.py:1422

  key, _ = _check_prng_key("dirichlet", key)
  dtype = dtypes.check_and_canonicalize_user_dtype(
      float if dtype is None else dtype)
  if not dtypes.issubdtype(dtype, np.floating):
    raise ValueError(f"dtype argument to `dirichlet` must be a float "
                     f"dtype, got {dtype}")
  if shape is not None:
    shape = core.canonicalize_shape(shape)
  out_sharding = canonicalize_sharding_for_samplers(out_sharding, "dirichlet", shape)
  return maybe_auto_axes(_dirichlet, out_sharding,
                         shape=shape, dtype=dtype)(key, alpha)

@jit(static_argnums=(2, 3))
def _dirichlet(key, alpha, shape, dtype) -> Array:
  from jax._src.nn.functions import softmax  # pyrefly: ignore[missing-import]

  if not np.ndim(alpha) >= 1:
    msg = "dirichlet requires alpha.ndim >= 1, got alpha.ndim == {}"
    raise ValueError(msg.format(np.ndim(alpha)))

  if shape is None:
    shape = np.shape(alpha)[:-1]
  else:
    _check_shape("dirichlet", shape, np.shape(alpha)[:-1])

  alpha = lax.convert_element_type(alpha, dtype)

  # Compute gamma in log space, otherwise small alpha can lead to poor behavior.
  log_gamma_samples = loggamma(key, alpha, shape + np.shape(alpha)[-1:], dtype)
  return softmax(log_gamma_samples, -1)


def exponential(key: ArrayLike,
                shape: Shape = (),
                dtype: DTypeLikeFloat | None = None,
                *,
                out_sharding: NamedSharding | P | None = None) -> Array:

View on GitHub (pinned to 1e1c6a8fc0)

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.

Example fix

// before
samples = jax.random.dirichlet(key, 2.0)  # scalar -> ValueError

// after
samples = jax.random.dirichlet(key, jnp.full((5,), 2.0))  # shape (5,) concentrations

Defensive patterns

Strategy: validation

Validate before calling

import numpy as np, jax.numpy as jnp
alpha = jnp.atleast_1d(jnp.asarray(alpha, dtype=float))
assert np.ndim(alpha) >= 1

Type guard

def is_valid_dirichlet_alpha(alpha) -> bool:
    return np.ndim(alpha) >= 1

Prevention

Always construct alpha as a 1+-dim array, e.g. jnp.full((k,), val).
Wrap external config inputs with jnp.atleast_1d before sampling.

When it happens

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

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

Related errors

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