jax-ml/jax · error · NotImplementedError

axes argument to transpose()

Error message

axes argument to transpose()

What it means

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.

Source

Thrown at jax/experimental/sparse/coo.py:150

    if diag_size <= 0:
      # if k is out of range, return an empty matrix.
      return cls._empty((N, M), dtype=dtype, index_dtype=index_dtype)

    data = jnp.ones(diag_size, dtype=dtype)
    idx = jnp.arange(diag_size, dtype=index_dtype)
    zero = _const(idx, 0)
    k = _const(idx, k)
    row = lax.sub(idx, lax.cond(k >= 0, lambda: zero, lambda: k))
    col = lax.add(idx, lax.cond(k <= 0, lambda: zero, lambda: k))
    return cls((data, row, col), shape=(N, M), rows_sorted=True, cols_sorted=True)

  def todense(self) -> Array:
    return coo_todense(self)

  def transpose(self, axes: tuple[int, ...] | None = None) -> COO:
    if axes is not None:
      raise NotImplementedError("axes argument to transpose()")
    return COO((self.data, self.col, self.row), shape=self.shape[::-1],
               rows_sorted=self._cols_sorted, cols_sorted=self._rows_sorted)

  def tree_flatten(self) -> tuple[tuple[Array, Array, Array], dict[str, Any]]:
    return (self.data, self.row, self.col), self._info._asdict()

  @classmethod
  def tree_unflatten(cls, aux_data, children):
    obj = object.__new__(cls)
    obj.data, obj.row, obj.col = children
    if aux_data.keys() != {'shape', 'rows_sorted', 'cols_sorted'}:
      raise ValueError(f"COO.tree_unflatten: invalid {aux_data=}")
    obj.shape = aux_data['shape']
    obj._rows_sorted = aux_data['rows_sorted']
    obj._cols_sorted = aux_data['cols_sorted']
    return obj

  def __matmul__(self, other: ArrayLike) -> Array:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. For 2D matrices call transpose() with no arguments (axes=None)
2. For general permutations, convert to BCOO and use its transpose, or todense() → transpose → back
3. Reorder axes at the dense level if the result is consumed densely anyway

Example fix

# before
mt = coo_mat.transpose(axes=(1, 0))  # NotImplementedError

# after
mt = coo_mat.transpose()  # plain 2D transpose
# or: mt = bcoo.BCOO.fromdense(coo_mat.todense()).transpose((1, 0))

assert axes is None or tuple(axes) == (1, 0) and len(axes) == 2, \
    'COO supports only plain 2D transpose'

try:
    mt = coo_mat.transpose(axes)
except NotImplementedError:
    mt = bcoo.BCOO.fromdense(coo_mat.todense()).transpose(tuple(axes))

When it happens

Trigger: 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.)

Common situations: Generic numeric code calling .transpose(axes) uniformly on array-likes; permuting batch dims of a sparse operand; porting dense code to sparse.

Related errors

• matmul between two sparse objects.
• for transpose support, subclass {type(self)} must implement
• Cannot permute last two dimensions with leading dimensions.
• Sparse metadata format not implemented for {operand_dtype=}
• Unsupported transforms: {a_sparse_metadata_transforms}

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