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.

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 with non-positive exponent is not…
• sparse rule for is not implemented.
• sparse rule for is not implemented because it would result…
• A must be ( , ) matrix A, got output

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

Appendix: 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)