jax-ml/jax · error · NotImplementedError

transpose permutation cannot permute batch axes with non-bat

Error message

transpose permutation cannot permute batch axes with non-batch axes; got permutation {permutation}, with {n_batch=}.

What it means

BCOO transpose cannot mix axis classes: batch axes may only permute among batch axes, sparse axes among sparse axes, dense among dense. A permutation that moves a batch axis into a sparse/dense slot (or vice versa) raises NotImplementedError.

Source

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

  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)
  data = data.transpose(*batch_perm, n_batch, *(d + n_batch + 1 for d in dense_perm))
  return data, indices

@bcoo_transpose_p.def_abstract_eval
def _bcoo_transpose_abstract_eval(data, indices, *, permutation: Sequence[int], spinfo: SparseInfo):
  batch_perm, _, dense_perm = _validate_permutation(data, indices, permutation, spinfo.shape)
  n_batch = len(batch_perm)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Split the permutation so it permutes batch dims, sparse dims, and dense dims separately (use swapaxes-style pairwise transposes)
2. Reorder axes by reconstructing: transpose batch dims and sparse dims with two valid permutations
3. Reshape via bcoo.reshape with shapes that keep batch/dense structure intact

Example fix

// before
out = bcoo.transpose((1, 0, 2))  # batch axis swapped with sparse axis
// after
out = bcoo.transpose((0, 2, 1))  # permute within allowed classes; adjust as needed

n_batch = bcoo.indices.ndim - 2
n_dense = len(bcoo.shape) - n_batch - bcoo.n_sparse
bp, sp_, dp = perm[:n_batch], perm[n_batch:n_batch+bcoo.n_sparse], perm[n_batch+bcoo.n_sparse:]
assert sorted(bp) == list(range(n_batch)) and sorted(dp) == list(range(n_dense)), 'cross-class permutation'

When it happens

Trigger: bcoo.transpose on a batched BCOO with a permutation like (1, 0) where axis 0 is batch and axis 1 is sparse; likewise for dense dims. Also triggered via bcoo_reshape paths.

Common situations: Applying a generic dense-style transpose permutation to a sparse/batched BCOO; reshaping batched sparse arrays in ways that cross axis classes.

Related errors

• transpose permutation isn't a permutation of operand dimensi
• batch_dims must be None or satisfy 0 < dim < n_batch. Got {b
• transpose of bcoo_extract with assume_unique=False
• bcoo_extract with unequal batch dimensions.
• {batch_dims=} out of range for indices with {n_batch=}

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