jax-ml/jax · error · ValueError

Number of samples without replacement ({k}) cannot exceed nu

Error message

Number of samples without replacement ({k}) cannot exceed number of categories ({logits_arr.shape[axis]}).

What it means

jax.random.categorical with an explicit output shape samples without replacement by taking the top-k Gumbel-perturbed logits along axis. k equals the product of the requested shape's leading dims (the number of samples drawn in parallel), so if you request more samples than there are categories along axis, top_k cannot fill the result and ValueError is raised.

Source

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

                         replace=replace, mode=mode)(key, logits_arr)

def _categorical(key, logits_arr, shape, batch_shape, axis, replace, mode) -> Array:
  shape_prefix = shape[:len(shape)-len(batch_shape)]
  if replace:
    if axis >= 0:
      axis -= len(logits_arr.shape)

    logits_shape = list(shape[len(shape) - len(batch_shape):])
    logits_shape.insert(axis % len(logits_arr.shape), logits_arr.shape[axis])
    return jnp.argmax(
        gumbel(key, (*shape_prefix, *logits_shape), logits_arr.dtype, mode=mode) +
        lax.expand_dims(logits_arr, tuple(range(len(shape_prefix)))),
        axis=axis)
  else:
    logits_arr += gumbel(key, logits_arr.shape, logits_arr.dtype, mode=mode)
    k = math.prod(shape_prefix)
    if k > logits_arr.shape[axis]:
      raise ValueError(
        f"Number of samples without replacement ({k}) cannot exceed number of "
        f"categories ({logits_arr.shape[axis]})."
      )

    _, indices = lax.top_k(jnp.moveaxis(logits_arr, axis, -1), k)
    assert indices.shape == batch_shape + (k,)
    assert shape == shape_prefix + batch_shape

    dimensions = (indices.ndim - 1, *range(indices.ndim - 1))
    indices = lax.reshape(indices, shape, dimensions)
    assert indices.shape == shape
    return indices


def laplace(key: ArrayLike,
            shape: Shape = (),
            dtype: DTypeLikeFloat | None = None,
            *,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Reduce the number of samples: keep math.prod(shape[:-1]) <= logits.shape[axis], e.g. shape=(4,) for 4 categories.
2. Or sample with replacement by drawing single samples in a loop/vmap over fresh keys: vmap(lambda k: jax.random.categorical(k, logits))(jax.random.split(key, n)).
3. Double-check the axis argument so k is compared against the true category dimension.

Example fix

// before
idx = jax.random.categorical(key, logits, shape=(6,))  # logits has 4 categories

// after
keys = jax.random.split(key, 6)
idx = jax.vmap(lambda k: jax.random.categorical(k, logits))(keys)  # with replacement

import math, numpy as np
k = math.prod(tuple(shape)[:-1])
assert k <= logits.shape[axis], 'more samples than categories without replacement'

def samples_fit_categories(logits, shape, axis=-1) -> bool:
    import math
    return math.prod(tuple(shape)[:-1]) <= logits.shape[axis]

When it happens

Trigger: jax.random.categorical(key, logits, shape=(6,)) with logits.shape[-1] == 4 (drawing 6 samples from 4 categories without replacement); any case where math.prod(shape[:-1]) > logits.shape[axis].

Common situations: Assuming categorical with a shape samples with replacement like a multinomial; increasing batch/sample counts during data-generation refactors without growing the logits axis; axis argument pointing at the wrong dimension so the 'category' axis is smaller than intended.

Related errors

• dirichlet requires alpha.ndim >= 1, got alpha.ndim == {}
• lam shape must be broadcastable to shape argument; got lam.s
• {name} was requested to map its argument along axis {axis},
• {''.join(msg)[:-2]}
• jvp called with different primal and tangent shapes;Got prim

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