{"record":{"id":"a603a34de8efb261","repo":"jax-ml/jax","slug":"number-of-samples-without-replacement-k-cannot","errorCode":null,"errorMessage":"Number of samples without replacement ({k}) cannot exceed number of categories ({logits_arr.shape[axis]}).","messagePattern":"Number of samples without replacement \\((.+?)\\) cannot exceed number of categories \\((.+?)\\)\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/random/core.py","lineNumber":2419,"sourceCode":"                         replace=replace, mode=mode)(key, logits_arr)\n\ndef _categorical(key, logits_arr, shape, batch_shape, axis, replace, mode) -> Array:\n  shape_prefix = shape[:len(shape)-len(batch_shape)]\n  if replace:\n    if axis >= 0:\n      axis -= len(logits_arr.shape)\n\n    logits_shape = list(shape[len(shape) - len(batch_shape):])\n    logits_shape.insert(axis % len(logits_arr.shape), logits_arr.shape[axis])\n    return jnp.argmax(\n        gumbel(key, (*shape_prefix, *logits_shape), logits_arr.dtype, mode=mode) +\n        lax.expand_dims(logits_arr, tuple(range(len(shape_prefix)))),\n        axis=axis)\n  else:\n    logits_arr += gumbel(key, logits_arr.shape, logits_arr.dtype, mode=mode)\n    k = math.prod(shape_prefix)\n    if k > logits_arr.shape[axis]:\n      raise ValueError(\n        f\"Number of samples without replacement ({k}) cannot exceed number of \"\n        f\"categories ({logits_arr.shape[axis]}).\"\n      )\n\n    _, indices = lax.top_k(jnp.moveaxis(logits_arr, axis, -1), k)\n    assert indices.shape == batch_shape + (k,)\n    assert shape == shape_prefix + batch_shape\n\n    dimensions = (indices.ndim - 1, *range(indices.ndim - 1))\n    indices = lax.reshape(indices, shape, dimensions)\n    assert indices.shape == shape\n    return indices\n\n\ndef laplace(key: ArrayLike,\n            shape: Shape = (),\n            dtype: DTypeLikeFloat | None = None,\n            *,","sourceCodeStart":2401,"sourceCodeEnd":2437,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/random/core.py#L2401-L2437","documentation":"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.","triggerScenarios":"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].","commonSituations":"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.","solutions":["Reduce the number of samples: keep math.prod(shape[:-1]) <= logits.shape[axis], e.g. shape=(4,) for 4 categories.","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)).","Double-check the axis argument so k is compared against the true category dimension."],"exampleFix":"// before\nidx = jax.random.categorical(key, logits, shape=(6,))  # logits has 4 categories\n\n// after\nkeys = jax.random.split(key, 6)\nidx = jax.vmap(lambda k: jax.random.categorical(k, logits))(keys)  # with replacement","handlingStrategy":"validation","validationCode":"import math, numpy as np\nk = math.prod(tuple(shape)[:-1])\nassert k <= logits.shape[axis], 'more samples than categories without replacement'","typeGuard":"def samples_fit_categories(logits, shape, axis=-1) -> bool:\n    import math\n    return math.prod(tuple(shape)[:-1]) <= logits.shape[axis]","tryCatchPattern":null,"preventionTips":["For with-replacement sampling use vmap over split keys, not a big shape.","Check logits.shape[axis] against requested sample count before calling."],"tags":["jax","random","categorical","sampling","top-k","shape"],"backgroundTag":"sample-count-exceeds-population","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}