{"record":{"id":"6b3ee63ab1155b2b","repo":"jax-ml/jax","slug":"todense-transpose-for-type-obj","errorCode":null,"errorMessage":"todense_transpose for {type(obj)}","messagePattern":"todense_transpose for (.+?)","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/experimental/sparse/api.py","lineNumber":107,"sourceCode":"\n  standin = object()\n  obj = tree_util.tree_unflatten(tree, [standin] * len(bufs))\n  from jax.experimental.sparse import BCOO, BCSR\n  from jax.experimental.sparse.bcoo import _bcoo_extract\n  from jax.experimental.sparse.bcsr import bcsr_extract\n  if obj is standin:\n    return (ct,)\n  elif isinstance(obj, BCOO):\n    _, indices = bufs\n    return _bcoo_extract(indices, ct), indices\n  elif isinstance(obj, BCSR):\n    _, indices, indptr = bufs\n    return bcsr_extract(indices, indptr, ct), indices, indptr\n  elif isinstance(obj, COO):\n    _, row, col = bufs\n    return _coo_extract(row, col, ct), row, col\n  else:\n    raise NotImplementedError(f\"todense_transpose for {type(obj)}\")\n\ndef _todense_batching_rule(batched_args, batch_dims, *, tree):\n  return jax.vmap(partial(_todense_impl, tree=tree), batch_dims)(*batched_args), 0\n\nad.primitive_jvps[todense_p] = _todense_jvp\nad.primitive_transposes[todense_p] = _todense_transpose\nbatching.primitive_batchers[todense_p] = _todense_batching_rule\nmlir.register_lowering(todense_p, mlir.lower_fun(\n    _todense_impl, multiple_results=False))\n\n\ndef empty(shape: Sequence[int], dtype: DTypeLike | None=None, index_dtype: DTypeLike = 'int32',\n          sparse_format: str = 'bcoo', **kwds) -> JAXSparse:\n  \"\"\"Create an empty sparse array.\n\n  Args:\n    shape: sequence of integers giving the array shape.\n    dtype: (optional) dtype of the array.","sourceCodeStart":89,"sourceCodeEnd":125,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/sparse/api.py#L89-L125","documentation":"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.","triggerScenarios":"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.","commonSituations":"Mixing the older CSR/CSC API with autodiff; migrating old code that used todense on CSR and later adding gradient computation.","solutions":["Convert the matrix to COO before calling todense: mat = mat.tocoo() (or BCSR with to_bcsr) so the transpose rule is defined","Restructure so the gradient does not flow through todense of a CSR/CSC (extract values via mat.data / sparsify ops instead)"],"exampleFix":"// before\ng = jax.grad(lambda m: f(jax.sparse.todense(m)))(csr_mat)\n// after\ng = jax.grad(lambda m: f(jax.sparse.todense(m)))(csr_mat.tocoo())","handlingStrategy":"validation","validationCode":"if not isinstance(mat, (jax.experimental.sparse.BCSR, jax.experimental.sparse.COO)):\n    mat = mat.tocoo() if hasattr(mat, 'tocoo') else mat.tobcoo()","typeGuard":"def is_transpose_supported_sparse(m) -> bool:\n    from jax.experimental.sparse import BCSR, COO\n    return isinstance(m, (BCSR, COO))","tryCatchPattern":"try:\n    jax.grad(f)(mat)\nexcept NotImplementedError:\n    jax.grad(f)(mat.tocoo())","preventionTips":["Standardize on BCOO when gradients are needed","Run a tiny smoke-test gradient at module import for new sparse types"],"tags":["jax","sparse","autodiff","todense","csr"],"backgroundTag":"unsupported-autodiff-rule","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}