{"record":{"id":"f7882534d86b3e83","repo":"jax-ml/jax","slug":"axes-argument-to-transpose","errorCode":null,"errorMessage":"axes argument to transpose()","messagePattern":"axes argument to transpose\\(\\)","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/experimental/sparse/coo.py","lineNumber":150,"sourceCode":"\n    if diag_size <= 0:\n      # if k is out of range, return an empty matrix.\n      return cls._empty((N, M), dtype=dtype, index_dtype=index_dtype)\n\n    data = jnp.ones(diag_size, dtype=dtype)\n    idx = jnp.arange(diag_size, dtype=index_dtype)\n    zero = _const(idx, 0)\n    k = _const(idx, k)\n    row = lax.sub(idx, lax.cond(k >= 0, lambda: zero, lambda: k))\n    col = lax.add(idx, lax.cond(k <= 0, lambda: zero, lambda: k))\n    return cls((data, row, col), shape=(N, M), rows_sorted=True, cols_sorted=True)\n\n  def todense(self) -> Array:\n    return coo_todense(self)\n\n  def transpose(self, axes: tuple[int, ...] | None = None) -> COO:\n    if axes is not None:\n      raise NotImplementedError(\"axes argument to transpose()\")\n    return COO((self.data, self.col, self.row), shape=self.shape[::-1],\n               rows_sorted=self._cols_sorted, cols_sorted=self._rows_sorted)\n\n  def tree_flatten(self) -> tuple[tuple[Array, Array, Array], dict[str, Any]]:\n    return (self.data, self.row, self.col), self._info._asdict()\n\n  @classmethod\n  def tree_unflatten(cls, aux_data, children):\n    obj = object.__new__(cls)\n    obj.data, obj.row, obj.col = children\n    if aux_data.keys() != {'shape', 'rows_sorted', 'cols_sorted'}:\n      raise ValueError(f\"COO.tree_unflatten: invalid {aux_data=}\")\n    obj.shape = aux_data['shape']\n    obj._rows_sorted = aux_data['rows_sorted']\n    obj._cols_sorted = aux_data['cols_sorted']\n    return obj\n\n  def __matmul__(self, other: ArrayLike) -> Array:","sourceCodeStart":132,"sourceCodeEnd":168,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/sparse/coo.py#L132-L168","documentation":"jax.experimental.sparse.COO.transpose only implements the plain 2D matrix transpose (swap row and col buffers, reverse shape). Arbitrary axis permutations are not implemented for this legacy format, so passing an axes tuple raises NotImplementedError. BCOO supports general transpose.","triggerScenarios":"coo_array.transpose(axes=(1,0)) or any non-None axes argument; also hit internally via _coo_todense_gpu_lowering paths that assume plain transposition. (For a 2D COO, transpose() with no args works fine.)","commonSituations":"Generic numeric code calling .transpose(axes) uniformly on array-likes; permuting batch dims of a sparse operand; porting dense code to sparse.","solutions":["For 2D matrices call transpose() with no arguments (axes=None)","For general permutations, convert to BCOO and use its transpose, or todense() → transpose → back","Reorder axes at the dense level if the result is consumed densely anyway"],"exampleFix":"# before\nmt = coo_mat.transpose(axes=(1, 0))  # NotImplementedError\n\n# after\nmt = coo_mat.transpose()  # plain 2D transpose\n# or: mt = bcoo.BCOO.fromdense(coo_mat.todense()).transpose((1, 0))","handlingStrategy":"fallback","validationCode":"assert axes is None or tuple(axes) == (1, 0) and len(axes) == 2, \\\n    'COO supports only plain 2D transpose'","typeGuard":null,"tryCatchPattern":"try:\n    mt = coo_mat.transpose(axes)\nexcept NotImplementedError:\n    mt = bcoo.BCOO.fromdense(coo_mat.todense()).transpose(tuple(axes))","preventionTips":["Call transpose() without args for 2D matrices","Use BCOO when general axis permutations are needed"],"tags":["jax","sparse","coo","transpose","not-implemented"],"backgroundTag":"sparse-op-unsupported","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}