{"record":{"id":"5ceced23e6b66b8f","repo":"jax-ml/jax","slug":"cannot-transpose-with-respect-to-sparse-indices-5ceced","errorCode":null,"errorMessage":"Cannot transpose with respect to sparse indices","messagePattern":"Cannot transpose with respect to sparse indices","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/sparse/csr.py","lineNumber":288,"sourceCode":"  dtype = data_aval.dtype\n  if not (np.issubdtype(dtype, np.floating) or np.issubdtype(dtype, np.complexfloating)):\n    warnings.warn(f\"csr_todense cusparse/hipsparse lowering not available for {dtype=}. \"\n                  \"Falling back to default implementation.\", CuSparseEfficiencyWarning)\n    return _csr_todense_lowering(ctx, data, indices, indptr, shape=shape)\n  return [_lowerings.csr_todense_gpu_lowering(\n      ctx, data, indices, indptr, shape=shape,\n      target_name_prefix=target_name_prefix)]\n\n\ndef _csr_todense_jvp(data_dot, data, indices, indptr, *, shape):\n  return _csr_todense(data_dot, indices, indptr, shape=shape)\n\ndef _csr_todense_transpose(ct, data, indices, indptr, *, shape):\n  # Note: we assume that transpose has the same sparsity pattern.\n  # Can we check this?\n  assert ad.is_undefined_primal(data)\n  if ad.is_undefined_primal(indices) or ad.is_undefined_primal(indptr):\n    raise ValueError(\"Cannot transpose with respect to sparse indices\")\n  assert ct.shape == shape\n  assert indices.aval.dtype == indptr.aval.dtype\n  assert ct.dtype == data.aval.dtype\n  return _csr_extract(indices, indptr, ct), indices, indptr\n\nad.defjvp(csr_todense_p, _csr_todense_jvp, None, None)\nad.primitive_transposes[csr_todense_p] = _csr_todense_transpose\nmlir.register_lowering(csr_todense_p, _csr_todense_lowering)\ndispatch.simple_impl(csr_todense_p)\n\nmlir.register_lowering(\n    csr_todense_p,\n    partial(_csr_todense_gpu_lowering, target_name_prefix='cu'),\n    platform='cuda')\nmlir.register_lowering(\n    csr_todense_p,\n    partial(_csr_todense_gpu_lowering, target_name_prefix='hip'),\n    platform='rocm')","sourceCodeStart":270,"sourceCodeEnd":306,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/sparse/csr.py#L270-L306","documentation":"During JAX's transpose (backward-pass) rule for csr_todense, the CSR index arrays (indices, indptr) themselves were traced as differentiable values. Gradients with respect to sparse structure are not defined/differentiable, so JAX raises.","triggerScenarios":"Calling jax.grad (or vjp/jvp backward) on a function whose csr_todense (e.g. sparse.CSR.dot, todense) has indices/indptr as traced/differentiated arguments, e.g. differentiating through code that produces index arrays from parameters.","commonSituations":"Wrapping sparse construction (jnp.argsort, searchsorted outputs feeding indices) inside a differentiated function; using lax.custom_sparse operations where indices depend on inputs.","solutions":["Use jax.lax.stop_gradient on indices/indptr before passing them into the sparse op","Keep index arrays as static/non-traced inputs (construct CSR outside grad scope)","Differentiate only with respect to the data buffer"],"exampleFix":"// before\nM = CSR((data, idx, indptr))\nout = jax.grad(lambda d: todense(CSR((d, idx, indptr))).sum())(data)  # idx traced\n// after\nidx = jax.lax.stop_gradient(idx)\nindptr = jax.lax.stop_gradient(indptr)\nout = jax.grad(lambda d: todense(CSR((d, idx, indptr))).sum())(data)","handlingStrategy":"validation","validationCode":"import jax\nassert not isinstance(indices, jax.core.Tracer), 'indices must not be traced'\nassert not isinstance(indptr, jax.core.Tracer), 'indptr must not be traced'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Treat CSR indices/indptr as static structure: never compute them from traced values","Wrap index arrays in jax.lax.stop_gradient when they might enter a grad trace","Only differentiate with respect to the data buffer"],"tags":["jax","sparse","csr","autodiff","gradient"],"backgroundTag":"gradient-through-sparse-indices","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}