{"record":{"id":"41d60d5de91a34de","repo":"jax-ml/jax","slug":"addition-between-sparse-matrices-of-different-shap","errorCode":null,"errorMessage":"Addition between sparse matrices of different shapes.","messagePattern":"Addition between sparse matrices of different shapes\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/experimental/sparse/transform.py","lineNumber":639,"sourceCode":"\nsparse_rules_bcoo[lax.stack_p] = functools.partial(\n    _stack_sparse,\n    broadcast_in_dim=sparse.bcoo_broadcast_in_dim,\n    concatenate=sparse.bcoo_concatenate,\n)\nsparse_rules_bcsr[lax.stack_p] = functools.partial(\n    _stack_sparse,\n    broadcast_in_dim=sparse.bcsr_broadcast_in_dim,\n    concatenate=sparse.bcsr_concatenate,\n)\n\n\ndef _add_sparse(spenv, *spvalues):\n  X, Y = spvalues\n  out_shape = lax.broadcast_shapes(X.shape, Y.shape)\n  if X.is_sparse() and Y.is_sparse():\n    if X.shape != Y.shape:\n      raise NotImplementedError(\"Addition between sparse matrices of different shapes.\")\n    if X.indices_ref == Y.indices_ref:\n      out_data = lax.add(spenv.data(X), spenv.data(Y))\n      if config.enable_checks.value:\n        assert X.indices_sorted == Y.indices_sorted\n        assert X.unique_indices == Y.unique_indices\n      out_spvalue = spenv.sparse(X.shape, out_data, indices_ref=X.indices_ref,\n                                 indices_sorted=X.indices_sorted,\n                                 unique_indices=X.unique_indices)\n    elif spenv.indices(X).ndim != spenv.indices(Y).ndim or spenv.data(X).ndim != spenv.data(Y).ndim:\n      raise NotImplementedError(\"Addition between sparse matrices with different batch/dense dimensions.\")\n    else:\n      out_indices = lax.concatenate([spenv.indices(X), spenv.indices(Y)], dimension=spenv.indices(X).ndim - 2)\n      out_data = lax.concatenate([spenv.data(X), spenv.data(Y)], dimension=spenv.indices(X).ndim - 2)\n      out_spvalue = spenv.sparse(X.shape, out_data, out_indices)\n  else:\n    if Y.is_sparse():\n      X, Y = Y, X\n    assert X.is_sparse() and Y.is_dense()","sourceCodeStart":621,"sourceCodeEnd":657,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/sparse/transform.py#L621-L657","documentation":"The sparse add rule in sparsify supports elementwise addition of two sparse values only when their shapes match exactly (broadcasting two sparse structures is not implemented). Mismatched shapes raise NotImplementedError.","triggerScenarios":"Inside a @sparse.sparsify function, adding two sparse BCOO values of different shapes, e.g. (4,4) + (4,1) or (4,4) + (5,5), where at least the broadcast result would need sparse structure synthesis.","commonSituations":"Adding a sparse matrix to a sparse row/column vector; combining sparse matrices of different sizes; intended numpy-style broadcasting between two sparse operands.","solutions":["Match shapes before adding: slice/pad one operand so both have identical shapes","Convert one operand to dense with sparse.todense (sparse+dense is supported)","Broadcast the smaller sparse operand's data/indices manually to the target shape"],"exampleFix":"// before\n@sparse.sparsify\ndef f(M, v):  # v sparse (4,1)\n  return M + v\n// after\n@sparse.sparsify\ndef f(M, v):\n  return M + sparse.todense(v)","handlingStrategy":"validation","validationCode":"assert X.shape == Y.shape or not (X.is_sparse() and Y.is_sparse()), \\\n    'sparse + sparse requires identical shapes'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Match shapes (slice/pad) before adding two sparse matrices","Convert one operand to dense for broadcasting semantics"],"tags":["jax","sparse","sparsify","addition","shape-mismatch"],"backgroundTag":"sparse-broadcasting-unsupported","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}