jax-ml/jax · error · ValueError

`perm` does not represent a permutation: {perm}

Error message

`perm` does not represent a permutation: {perm}

What it means

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.

Source

Thrown at jax/_src/lax/parallel.py:512

  If ``x`` is a pytree then the result is equivalent to mapping this function to
  each leaf in the tree.

  Args:
    x: array(s) with a mapped axis named ``axis_name``.
    axis_name: hashable Python object used to name a pmapped axis (see the
      :func:`jax.pmap` documentation for more details).
    perm: list of ints encoding sources for the permutation to be applied to
      the axis named ``axis_name``, so that the output at axis index i
      comes from the input at axis index perm[i]. Every integer in [0, N) should
      be included exactly once for axis size N.

  Returns:
    Array(s) with the same shape as ``x`` with slices along the axis
    ``axis_name`` gathered from ``x`` according to the permutation ``perm``.
  """
  if set(perm) != set(range(len(perm))):
    raise ValueError(f"`perm` does not represent a permutation: {perm}")
  return ppermute(x, axis_name, list(zip(perm, range(len(perm)))))


def pswapaxes(x, axis_name, axis, *, axis_index_groups=None):
  """Swap the pmapped axis ``axis_name`` with the unmapped axis ``axis``.

  If ``x`` is a pytree then the result is equivalent to mapping this function to
  each leaf in the tree.

  The group size of the mapped axis size must be equal to the size of the
  unmapped axis; that is, we must have
  ``lax.psum(1, axis_name, axis_index_groups=axis_index_groups) == x.shape[axis]``.
  By default, when ``axis_index_groups=None``, this encompasses all the devices.

  This function is a special case of ``all_to_all`` where the pmapped axis of
  the input is placed at the position ``axis`` in the output. That is, it is
  equivalent to ``all_to_all(x, axis_name, axis, axis)``.

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Validate perm covers 0..n-1 exactly once: sorted(perm) == list(range(len(perm)))
2. If devices may legitimately be dropped, use ppermute with explicit (src,dst) pairs for only the moved entries
3. Regenerate perm when the axis size changes

Example fix

# before
y = jax.lax.pshuffle(x, 'i', [0, 0, 2])

# after
y = jax.lax.pshuffle(x, 'i', [0, 2, 1])

assert sorted(perm) == list(range(len(perm))), f'not a permutation: {perm}'

def is_permutation(perm):
    return sorted(perm) == list(range(len(perm)))

When it happens

Trigger: pshuffle(x, 'i', [0,0,2]) (duplicate 0, missing 1) or [0,1,3] (out of range for length 3).

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

Related errors

• {full_name} must be a pytree prefix with bool leaves or a tu
• Axes mentioned in `manual_axis_type` field of ShapedArray sh
• varying and unreduced cannot have common mesh axes. Got vary
• varying and reduced cannot have common mesh axes. Got varyin
• `unreduced_kind` should be `None` when `unreduced` is an emp

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