jax-ml/jax · error · NotImplementedError

Addition between sparse matrices of different shapes.

Error message

Addition between sparse matrices of different shapes.

What it means

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.

Source

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

sparse_rules_bcoo[lax.stack_p] = functools.partial(
    _stack_sparse,
    broadcast_in_dim=sparse.bcoo_broadcast_in_dim,
    concatenate=sparse.bcoo_concatenate,
)
sparse_rules_bcsr[lax.stack_p] = functools.partial(
    _stack_sparse,
    broadcast_in_dim=sparse.bcsr_broadcast_in_dim,
    concatenate=sparse.bcsr_concatenate,
)


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()

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Match shapes before adding: slice/pad one operand so both have identical shapes
2. Convert one operand to dense with sparse.todense (sparse+dense is supported)
3. Broadcast the smaller sparse operand's data/indices manually to the target shape

Example fix

// before
@sparse.sparsify
def f(M, v):  # v sparse (4,1)
  return M + v
// after
@sparse.sparsify
def f(M, v):
  return M + sparse.todense(v)

assert X.shape == Y.shape or not (X.is_sparse() and Y.is_sparse()), \
    'sparse + sparse requires identical shapes'

When it happens

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

Common situations: Adding a sparse matrix to a sparse row/column vector; combining sparse matrices of different sizes; intended numpy-style broadcasting between two sparse operands.

Related errors

• Addition between sparse matrices with different batch/dense
• bcoo_slice: indices must have size mat.ndim={mat.ndim}
• bcoo_dynamic_slice: indices must have size mat.ndim={mat.ndi
• bcoo_multiply_sparse: arrays must have same number of dimens
• A must be ({n}, {n}) matrix A, got output {s}

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