jax-ml/jax · error · TypeError

`perm` passed to `jax.lax.ppermute` must be a list or a tupl

Error message

`perm` passed to `jax.lax.ppermute` must be a list or a tuple. Got perm of type {type(perm)}

What it means

jax.lax.ppermute requires perm to be a list or tuple of (source, destination) index pairs describing a permutation across the mapped axis. Passing e.g. a dict, numpy array, or generator raises TypeError before any device communication.

Source

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

      ``(source_index, destination_index)``
      pairs that encode how the mapped axis named ``axis_name`` should be
      shuffled. The integer values are treated as indices into the mapped axis
      ``axis_name``. Any two pairs should not have the same source index or the
      same destination index. For each index of the axis ``axis_name`` that does
      not correspond to a destination index in ``perm``, the corresponding
      values in the result are filled with zeros of the appropriate type.

  Returns:
    Array(s) with the same shape as ``x`` with slices along the axis
    ``axis_name`` gathered from ``x`` according to the permutation ``perm``.
  """
  return _ppermute_is_async(x, axis_name, perm, is_async=False)

def _ppermute_is_async(x, axis_name, perm, is_async=False):
  if not isinstance(axis_name, (list, tuple)):
    axis_name = (axis_name,)
  if not isinstance(perm, (list, tuple)):
    raise TypeError(
        "`perm` passed to `jax.lax.ppermute` must be a list or a tuple. Got"
        f" perm of type {type(perm)}")
  def bind(leaf):
    leaf = insert_collective_pvary(axis_name, leaf)
    prim = ppermute_start_p if is_async else ppermute_p
    return prim.bind(leaf, axis_name=axis_name, perm=tuple(map(tuple, perm)))
  return tree_util.tree_map(bind, x)


def psend(x, axis_name, perm):
  """Perform a collective send according to the permutation ``perm``.

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

  This function is an analog of the Send HLO.

  Args:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Convert to a list/tuple of 2-tuples: perm = [(int(a), int(b)) for a, b in perm]
2. If using pshuffle semantics (flat permutation), call pshuffle(x, axis_name, perm) which builds the inverse pairs for you

Example fix

// before
y = jax.lax.ppermute(x, 'i', np.array([[0,1],[1,0]]))

// after
y = jax.lax.ppermute(x, 'i', [(0,1),(1,0)])

assert isinstance(perm, (list, tuple)) and all(isinstance(p, (list, tuple)) and len(p) == 2 for p in perm)

def valid_ppermute_perm(perm):
    return isinstance(perm, (list, tuple)) and all(
        isinstance(p, (list, tuple)) and len(p) == 2 for p in perm)

When it happens

Trigger: ppermute(x, 'i', np.array([[0,1],[1,0]])) or perm as a dict/generator/iterator.

Common situations: Passing a numpy array of pairs from routing logic; converting perm data from another format and forgetting to list()-ify.

Related errors

• {op_name} sources and destinations must be unique, got {}.
• iteration over a 0-d array
• Formatting arguments to checkify.check need to be PyTrees of
• check_error takes an Error as argument, got type {type(error
• bool() not supported for instances of type '{0}' (did you me

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