jax-ml/jax · error · NotImplementedError

todense_transpose for {type(obj)}

Error message

todense_transpose for {type(obj)}

What it means

The transpose rule for sparse todense only knows how to pull cotangents back through BCSR and COO objects. Differentiating a todense call on any other sparse type (e.g. CSR/CSC) raises NotImplementedError.

Source

Thrown at jax/experimental/sparse/api.py:107

  standin = object()
  obj = tree_util.tree_unflatten(tree, [standin] * len(bufs))
  from jax.experimental.sparse import BCOO, BCSR
  from jax.experimental.sparse.bcoo import _bcoo_extract
  from jax.experimental.sparse.bcsr import bcsr_extract
  if obj is standin:
    return (ct,)
  elif isinstance(obj, BCOO):
    _, indices = bufs
    return _bcoo_extract(indices, ct), indices
  elif isinstance(obj, BCSR):
    _, indices, indptr = bufs
    return bcsr_extract(indices, indptr, ct), indices, indptr
  elif isinstance(obj, COO):
    _, row, col = bufs
    return _coo_extract(row, col, ct), row, col
  else:
    raise NotImplementedError(f"todense_transpose for {type(obj)}")

def _todense_batching_rule(batched_args, batch_dims, *, tree):
  return jax.vmap(partial(_todense_impl, tree=tree), batch_dims)(*batched_args), 0

ad.primitive_jvps[todense_p] = _todense_jvp
ad.primitive_transposes[todense_p] = _todense_transpose
batching.primitive_batchers[todense_p] = _todense_batching_rule
mlir.register_lowering(todense_p, mlir.lower_fun(
    _todense_impl, multiple_results=False))


def empty(shape: Sequence[int], dtype: DTypeLike | None=None, index_dtype: DTypeLike = 'int32',
          sparse_format: str = 'bcoo', **kwds) -> JAXSparse:
  """Create an empty sparse array.

  Args:
    shape: sequence of integers giving the array shape.
    dtype: (optional) dtype of the array.

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Convert the matrix to COO before calling todense: mat = mat.tocoo() (or BCSR with to_bcsr) so the transpose rule is defined
2. Restructure so the gradient does not flow through todense of a CSR/CSC (extract values via mat.data / sparsify ops instead)

Example fix

// before
g = jax.grad(lambda m: f(jax.sparse.todense(m)))(csr_mat)
// after
g = jax.grad(lambda m: f(jax.sparse.todense(m)))(csr_mat.tocoo())

if not isinstance(mat, (jax.experimental.sparse.BCSR, jax.experimental.sparse.COO)):
    mat = mat.tocoo() if hasattr(mat, 'tocoo') else mat.tobcoo()

def is_transpose_supported_sparse(m) -> bool:
    from jax.experimental.sparse import BCSR, COO
    return isinstance(m, (BCSR, COO))

try:
    jax.grad(f)(mat)
except NotImplementedError:
    jax.grad(f)(mat.tocoo())

When it happens

Trigger: Computing gradients (jax.grad / jax.vjp / jax.jacfwd) through sparse.todense on a CSR or CSC matrix rather than BCSR/COO, hitting _todense_transpose with an unsupported obj type.

Common situations: Mixing the older CSR/CSC API with autodiff; migrating old code that used todense on CSR and later adding gradient computation.

Related errors

• Cannot transpose with respect to sparse indices
• transpose of bcoo_extract with assume_unique=False
• CSR must have ndim=2; got {shape=}
• matmul between two sparse objects.
• matmul with object of shape {other.shape}

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