jax-ml/jax · error · NotImplementedError

Addition between sparse matrices with different batch/dense

Error message

Addition between sparse matrices with different batch/dense dimensions.

What it means

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.

Source

Thrown at jax/experimental/sparse/transform.py:649

)


def _add_sparse(spenv, *spvalues):
  X, Y = spvalues
  out_shape = lax.broadcast_shapes(X.shape, Y.shape)
  if X.is_sparse() and Y.is_sparse():
    if X.shape != Y.shape:
      raise NotImplementedError("Addition between sparse matrices of different shapes.")
    if X.indices_ref == Y.indices_ref:
      out_data = lax.add(spenv.data(X), spenv.data(Y))
      if config.enable_checks.value:
        assert X.indices_sorted == Y.indices_sorted
        assert X.unique_indices == Y.unique_indices
      out_spvalue = spenv.sparse(X.shape, out_data, indices_ref=X.indices_ref,
                                 indices_sorted=X.indices_sorted,
                                 unique_indices=X.unique_indices)
    elif spenv.indices(X).ndim != spenv.indices(Y).ndim or spenv.data(X).ndim != spenv.data(Y).ndim:
      raise NotImplementedError("Addition between sparse matrices with different batch/dense dimensions.")
    else:
      out_indices = lax.concatenate([spenv.indices(X), spenv.indices(Y)], dimension=spenv.indices(X).ndim - 2)
      out_data = lax.concatenate([spenv.data(X), spenv.data(Y)], dimension=spenv.indices(X).ndim - 2)
      out_spvalue = spenv.sparse(X.shape, out_data, out_indices)
  else:
    if Y.is_sparse():
      X, Y = Y, X
    assert X.is_sparse() and Y.is_dense()
    if Y.shape != out_shape:
      raise NotImplementedError(
        "Addition between a sparse array X and a dense array Y is not implemented when "
        "the output shape is larger than Y.shape. This is to prevent silent densification "
        "of a large sparse array. If this is your intent, you can explicitly cast the sparse "
        "array to a dense matrix.")
    X_promoted, Y_promoted = spvalues_to_arrays(spenv, (X, Y))
    out = X_promoted.todense() + Y_promoted
    out_spvalue = spenv.dense(out)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Rebuild both operands with the same n_batch/n_dense layout (e.g. via sparse.bcoo.reshape or reconstructing with matching batch dims)
2. Convert one operand to dense before the add
3. Broadcast the unbatched operand's data to add a leading batch dim

Example fix

// before
@sparse.sparsify
def f(Mb, M):  # Mb batched, M not
  return Mb + M
// after
M_b = sparse.BCOO((jnp.broadcast_to(M.data, Mb.data.shape),
                   jnp.broadcast_to(M.indices, Mb.indices.shape)), shape=Mb.shape)
@sparse.sparsify
def f(Mb, M_b):
  return Mb + M_b

assert X.indices.ndim == Y.indices.ndim and X.data.ndim == Y.data.ndim, \
    'operands must share batch/dense layout'

When it happens

Trigger: 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.

Common situations: 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.

Related errors

• Addition between sparse matrices of different shapes.
• sparse rule for {primitive} is not implemented because it wo
• sparse rule for {primitive} is not implemented.
• sparse rule for {lax.integer_pow_p} with non-positive expone
• dot_general requires equal numbers of lhs_batch and rhs_batc

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