{"record":{"id":"f2672216e7e817c2","repo":"jax-ml/jax","slug":"transpose-permutation-cannot-permute-batch-axes-wi","errorCode":null,"errorMessage":"transpose permutation cannot permute batch axes with non-batch axes; got permutation {permutation}, with {n_batch=}.","messagePattern":"transpose permutation cannot permute batch axes with non-batch axes; got permutation (.+?), with (.+?)\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/experimental/sparse/bcoo.py","lineNumber":546,"sourceCode":"  permutation = tuple(permutation)\n  if permutation == tuple(range(len(spinfo.shape))):\n    return data, indices\n  else:\n    return bcoo_transpose_p.bind(data, indices, permutation=permutation,\n                                 spinfo=spinfo)\n\ndef _validate_permutation(data, indices, permutation, shape):\n  if not isinstance(permutation, (tuple, list, np.ndarray)):\n    raise TypeError(f\"transpose permutation must be a tuple/list/ndarray, got {type(permutation)}.\")\n  if tuple(sorted(permutation)) != tuple(range(len(shape))):\n    raise TypeError(\"transpose permutation isn't a permutation of operand dimensions, \"\n                    f\"got permutation {permutation} for shape {shape}.\")\n  n_batch, n_sparse, n_dense, _ = _validate_bcoo(data, indices, shape)\n  batch_perm = permutation[:n_batch]\n  sparse_perm = [p - n_batch for p in permutation[n_batch: n_batch + n_sparse]]\n  dense_perm = [p - n_sparse - n_batch for p in permutation[n_batch + n_sparse:]]\n  if n_batch and tuple(sorted(batch_perm)) != tuple(range(n_batch)):\n    raise NotImplementedError(\"transpose permutation cannot permute batch axes with non-batch axes; \"\n                              f\"got permutation {permutation}, with {n_batch=}.\")\n  if n_dense and tuple(sorted(dense_perm)) != tuple(range(n_dense)):\n    raise NotImplementedError(\"transpose permutation cannot permute dense axes with non-dense axes; \"\n                              f\"got permutation {permutation}, with {n_dense=}.\")\n  return batch_perm, sparse_perm, dense_perm\n\n@bcoo_transpose_p.def_impl\ndef _bcoo_transpose_impl(data, indices, *, permutation: Sequence[int], spinfo: SparseInfo):\n  batch_perm, sparse_perm, dense_perm = _validate_permutation(data, indices, permutation, spinfo.shape)\n  n_batch = len(batch_perm)\n  indices = indices[..., sparse_perm].transpose(*batch_perm, n_batch, n_batch + 1)\n  data = data.transpose(*batch_perm, n_batch, *(d + n_batch + 1 for d in dense_perm))\n  return data, indices\n\n@bcoo_transpose_p.def_abstract_eval\ndef _bcoo_transpose_abstract_eval(data, indices, *, permutation: Sequence[int], spinfo: SparseInfo):\n  batch_perm, _, dense_perm = _validate_permutation(data, indices, permutation, spinfo.shape)\n  n_batch = len(batch_perm)","sourceCodeStart":528,"sourceCodeEnd":564,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/sparse/bcoo.py#L528-L564","documentation":"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.","triggerScenarios":"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.","commonSituations":"Applying a generic dense-style transpose permutation to a sparse/batched BCOO; reshaping batched sparse arrays in ways that cross axis classes.","solutions":["Split the permutation so it permutes batch dims, sparse dims, and dense dims separately (use swapaxes-style pairwise transposes)","Reorder axes by reconstructing: transpose batch dims and sparse dims with two valid permutations","Reshape via bcoo.reshape with shapes that keep batch/dense structure intact"],"exampleFix":"// before\nout = bcoo.transpose((1, 0, 2))  # batch axis swapped with sparse axis\n// after\nout = bcoo.transpose((0, 2, 1))  # permute within allowed classes; adjust as needed","handlingStrategy":"validation","validationCode":"n_batch = bcoo.indices.ndim - 2\nn_dense = len(bcoo.shape) - n_batch - bcoo.n_sparse\nbp, sp_, dp = perm[:n_batch], perm[n_batch:n_batch+bcoo.n_sparse], perm[n_batch+bcoo.n_sparse:]\nassert sorted(bp) == list(range(n_batch)) and sorted(dp) == list(range(n_dense)), 'cross-class permutation'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Permute batch, sparse, and dense axes only within their own groups","Prefer bcoo.reshape or axis-wise swapaxes over general transpose for batched BCOO"],"tags":["jax","sparse","bcoo","transpose","permutation","batching"],"backgroundTag":"unsupported-axis-permutation","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}