{"record":{"id":"1767d0894322e4da","repo":"jax-ml/jax","slug":"addition-between-sparse-matrices-with-different-ba","errorCode":null,"errorMessage":"Addition between sparse matrices with different batch/dense dimensions.","messagePattern":"Addition between sparse matrices with different batch/dense dimensions\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/experimental/sparse/transform.py","lineNumber":649,"sourceCode":")\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()\n    if Y.shape != out_shape:\n      raise NotImplementedError(\n        \"Addition between a sparse array X and a dense array Y is not implemented when \"\n        \"the output shape is larger than Y.shape. This is to prevent silent densification \"\n        \"of a large sparse array. If this is your intent, you can explicitly cast the sparse \"\n        \"array to a dense matrix.\")\n    X_promoted, Y_promoted = spvalues_to_arrays(spenv, (X, Y))\n    out = X_promoted.todense() + Y_promoted\n    out_spvalue = spenv.dense(out)\n","sourceCodeStart":631,"sourceCodeEnd":667,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/sparse/transform.py#L631-L667","documentation":"Sparse addition falls back to concatenating index buffers when the operands don't share an indices_ref; this path requires the index arrays and data arrays to have matching rank (same number of batch and dense dims). Mismatched ndim between operands raises NotImplementedError.","triggerScenarios":"Adding two sparse values with different batch or dense dimension counts, e.g. a BCOO with n_batch=1 plus one with n_batch=0, inside a sparsify function.","commonSituations":"Combining sparse matrices built with different n_batch/n_dense layouts; adding a batched sparse matrix and an unbatched one; mixing BCOO and BCSR-style layouts.","solutions":["Rebuild both operands with the same n_batch/n_dense layout (e.g. via sparse.bcoo.reshape or reconstructing with matching batch dims)","Convert one operand to dense before the add","Broadcast the unbatched operand's data to add a leading batch dim"],"exampleFix":"// before\n@sparse.sparsify\ndef f(Mb, M):  # Mb batched, M not\n  return Mb + M\n// after\nM_b = sparse.BCOO((jnp.broadcast_to(M.data, Mb.data.shape),\n                   jnp.broadcast_to(M.indices, Mb.indices.shape)), shape=Mb.shape)\n@sparse.sparsify\ndef f(Mb, M_b):\n  return Mb + M_b","handlingStrategy":"validation","validationCode":"assert X.indices.ndim == Y.indices.ndim and X.data.ndim == Y.data.ndim, \\\n    'operands must share batch/dense layout'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Build sparse operands with identical n_batch/n_dense layouts","Broadcast unbatched sparse operands to the batched layout before adding"],"tags":["jax","sparse","sparsify","addition","batch-dims"],"backgroundTag":"sparse-broadcasting-unsupported","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}