{"record":{"id":"79c83cf2c2a1b2dc","repo":"jax-ml/jax","slug":"perm-does-not-represent-a-permutation-perm","errorCode":null,"errorMessage":"`perm` does not represent a permutation: {perm}","messagePattern":"`perm` does not represent a permutation: (.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/parallel.py","lineNumber":512,"sourceCode":"\n  If ``x`` is a pytree then the result is equivalent to mapping this function to\n  each leaf in the tree.\n\n  Args:\n    x: array(s) with a mapped axis named ``axis_name``.\n    axis_name: hashable Python object used to name a pmapped axis (see the\n      :func:`jax.pmap` documentation for more details).\n    perm: list of ints encoding sources for the permutation to be applied to\n      the axis named ``axis_name``, so that the output at axis index i\n      comes from the input at axis index perm[i]. Every integer in [0, N) should\n      be included exactly once for axis size N.\n\n  Returns:\n    Array(s) with the same shape as ``x`` with slices along the axis\n    ``axis_name`` gathered from ``x`` according to the permutation ``perm``.\n  \"\"\"\n  if set(perm) != set(range(len(perm))):\n    raise ValueError(f\"`perm` does not represent a permutation: {perm}\")\n  return ppermute(x, axis_name, list(zip(perm, range(len(perm)))))\n\n\ndef pswapaxes(x, axis_name, axis, *, axis_index_groups=None):\n  \"\"\"Swap the pmapped axis ``axis_name`` with the unmapped axis ``axis``.\n\n  If ``x`` is a pytree then the result is equivalent to mapping this function to\n  each leaf in the tree.\n\n  The group size of the mapped axis size must be equal to the size of the\n  unmapped axis; that is, we must have\n  ``lax.psum(1, axis_name, axis_index_groups=axis_index_groups) == x.shape[axis]``.\n  By default, when ``axis_index_groups=None``, this encompasses all the devices.\n\n  This function is a special case of ``all_to_all`` where the pmapped axis of\n  the input is placed at the position ``axis`` in the output. That is, it is\n  equivalent to ``all_to_all(x, axis_name, axis, axis)``.\n","sourceCodeStart":494,"sourceCodeEnd":530,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/parallel.py#L494-L530","documentation":"jax.lax.pshuffle derives a ppermute from a flat permutation perm; it requires perm to be a true permutation of range(len(perm)) — every index appearing exactly once. Duplicates, missing values, or out-of-range entries raise ValueError.","triggerScenarios":"pshuffle(x, 'i', [0,0,2]) (duplicate 0, missing 1) or [0,1,3] (out of range for length 3).","commonSituations":"Routing permutations from model logic that can be non-bijective (dropped devices); off-by-one index math; stale perm computed for a different axis size.","solutions":["Validate perm covers 0..n-1 exactly once: sorted(perm) == list(range(len(perm)))","If devices may legitimately be dropped, use ppermute with explicit (src,dst) pairs for only the moved entries","Regenerate perm when the axis size changes"],"exampleFix":"# before\ny = jax.lax.pshuffle(x, 'i', [0, 0, 2])\n\n# after\ny = jax.lax.pshuffle(x, 'i', [0, 2, 1])","handlingStrategy":"validation","validationCode":"assert sorted(perm) == list(range(len(perm))), f'not a permutation: {perm}'","typeGuard":"def is_permutation(perm):\n    return sorted(perm) == list(range(len(perm)))","tryCatchPattern":null,"preventionTips":["Unit-test routing permutations for bijectivity","Regenerate perms when device count changes"],"tags":["jax","pshuffle","permutation","validation"],"backgroundTag":"invalid-permutation","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}