jax-ml/jax · error · TypeError

transpose permutation isn't a permutation of operand dimensi

Error message

transpose permutation isn't a permutation of operand dimensions, got permutation {permutation} for shape {shape}.

What it means

The transpose permutation must be a valid permutation of range(len(shape)) — each axis index exactly once. Duplicates, missing axes, out-of-range values, or wrong length raise TypeError.

Source

Thrown at jax/experimental/sparse/bcoo.py:539

  """
  buffers = _bcoo_transpose(mat.data, mat.indices, permutation=permutation, spinfo=mat._info)
  out_shape = tuple(mat.shape[p] for p in permutation)
  return BCOO(buffers, shape=out_shape, unique_indices=mat.unique_indices)

def _bcoo_transpose(data: Array, indices: Array, *,
                    permutation: Sequence[int], spinfo: SparseInfo) -> tuple[Array, Array]:
  permutation = tuple(permutation)
  if permutation == tuple(range(len(spinfo.shape))):
    return data, indices
  else:
    return bcoo_transpose_p.bind(data, indices, permutation=permutation,
                                 spinfo=spinfo)

def _validate_permutation(data, indices, permutation, shape):
  if not isinstance(permutation, (tuple, list, np.ndarray)):
    raise TypeError(f"transpose permutation must be a tuple/list/ndarray, got {type(permutation)}.")
  if tuple(sorted(permutation)) != tuple(range(len(shape))):
    raise TypeError("transpose permutation isn't a permutation of operand dimensions, "
                    f"got permutation {permutation} for shape {shape}.")
  n_batch, n_sparse, n_dense, _ = _validate_bcoo(data, indices, shape)
  batch_perm = permutation[:n_batch]
  sparse_perm = [p - n_batch for p in permutation[n_batch: n_batch + n_sparse]]
  dense_perm = [p - n_sparse - n_batch for p in permutation[n_batch + n_sparse:]]
  if n_batch and tuple(sorted(batch_perm)) != tuple(range(n_batch)):
    raise NotImplementedError("transpose permutation cannot permute batch axes with non-batch axes; "
                              f"got permutation {permutation}, with {n_batch=}.")
  if n_dense and tuple(sorted(dense_perm)) != tuple(range(n_dense)):
    raise NotImplementedError("transpose permutation cannot permute dense axes with non-dense axes; "
                              f"got permutation {permutation}, with {n_dense=}.")
  return batch_perm, sparse_perm, dense_perm

@bcoo_transpose_p.def_impl
def _bcoo_transpose_impl(data, indices, *, permutation: Sequence[int], spinfo: SparseInfo):
  batch_perm, sparse_perm, dense_perm = _validate_permutation(data, indices, permutation, spinfo.shape)
  n_batch = len(batch_perm)
  indices = indices[..., sparse_perm].transpose(*batch_perm, n_batch, n_batch + 1)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Validate perm covers all axes: check sorted(perm) == list(range(ndim)) before calling
  2. Fix permutation generation logic for the actual rank

Example fix

// before
out = bcoo.transpose((0,))  # 2D matrix, missing axis 1
// after
out = bcoo.transpose((1, 0))
Defensive patterns

Strategy: validation

Validate before calling

assert sorted(map(int, perm)) == list(range(len(shape))), 'not a valid permutation'

Prevention

When it happens

Trigger: bcoo.transpose((0, 0)) on a 2D matrix, transpose((1, 2)) on a 2D shape, or a permutation with wrong length from reshape logic.

Common situations: Programmatically built permutations with bugs; reusing a permutation computed for a different-rank array; off-by-one axis indices.

Related errors


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