{"record":{"id":"d1bf1af32af252c4","repo":"jax-ml/jax","slug":"bcoo-multiply-sparse-arrays-with-differing-number","errorCode":null,"errorMessage":"bcoo_multiply_sparse: arrays with differing numbers of dense dimensions: {lhs}, {rhs}","messagePattern":"bcoo_multiply_sparse: arrays with differing numbers of dense dimensions: (.+?), (.+?)","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/experimental/sparse/bcoo.py","lineNumber":2219,"sourceCode":"  \"\"\"\n  out_data, out_indices, out_shape = _bcoo_multiply_sparse(\n      lhs.data, lhs.indices, rhs.data, rhs.indices, lhs_spinfo=lhs._info,\n      rhs_spinfo=rhs._info)\n  return BCOO((out_data, out_indices), shape=out_shape)\n\ndef _bcoo_multiply_sparse(lhs_data: Array, lhs_indices: Array, rhs_data: Array, rhs_indices: Array, *,\n                          lhs_spinfo: SparseInfo, rhs_spinfo: SparseInfo) -> tuple[Array, Array, Shape]:\n  lhs_shape = lhs_spinfo.shape\n  rhs_shape = rhs_spinfo.shape\n\n  lhs = _validate_bcoo(lhs_data, lhs_indices, lhs_shape)\n  rhs = _validate_bcoo(rhs_data, rhs_indices, rhs_shape)\n  if len(lhs_shape) != len(rhs_shape):\n    # Similar requirement as lax.mul:\n    raise TypeError(\"bcoo_multiply_sparse: arrays must have same number of dimensions, \"\n                    f\"got {lhs_shape}, {rhs_shape}\")\n  if lhs.n_dense != rhs.n_dense:\n    raise NotImplementedError(\"bcoo_multiply_sparse: arrays with differing numbers of \"\n                              f\"dense dimensions: {lhs}, {rhs}\")\n  n_batch = min(lhs.n_batch, rhs.n_batch)\n  _mul = functools.partial(_bcoo_multiply_sparse_unbatched,\n                           lhs_shape=lhs_shape[n_batch:],\n                           rhs_shape=rhs_shape[n_batch:])\n  _mul = nfold_vmap(_mul, n_batch)\n  data, indices = _mul(lhs_data, lhs_indices, rhs_data, rhs_indices)\n  return data, indices, jnp.broadcast_shapes(lhs_shape, rhs_shape)\n\ndef _bcoo_multiply_sparse_unbatched(lhs_data, lhs_indices, rhs_data, rhs_indices, *, lhs_shape, rhs_shape):\n  lhs = _validate_bcoo(lhs_data, lhs_indices, lhs_shape)\n  rhs = _validate_bcoo(rhs_data, rhs_indices, rhs_shape)\n  assert (lhs.n_batch == 0) or (rhs.n_batch == 0)  # Ensured at call site above\n\n  # TODO(jakevdp): this can be made more efficient by utilizing batch structure.\n  if lhs.n_batch:\n    lhs_data, lhs_indices = bcoo_update_layout(BCOO((lhs_data, lhs_indices), shape=lhs_shape), n_batch=0)._bufs\n    lhs = _validate_bcoo(lhs_data, lhs_indices, lhs_shape)","sourceCodeStart":2201,"sourceCodeEnd":2237,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/sparse/bcoo.py#L2201-L2237","documentation":"bcoo_multiply_sparse supports differing numbers of batch dimensions (it takes the min via vmap) but not differing numbers of dense dimensions; if lhs.n_dense != rhs.n_dense it raises NotImplementedError because no layout conversion is performed.","triggerScenarios":"Multiplying two BCOO arrays where one stores trailing dense dims (e.g. a block/ragged tensor with n_dense=1) and the other is a pure sparse matrix with n_dense=0.","commonSituations":"Mixing BCOO tensors created with different n_dense conventions (e.g. from _bcoo_fromdense with different dense_dimensions args); upgrading code that previously operated on uniform formats; multiplying a per-nonzero-vector tensor by a scalar-per-nonzero matrix.","solutions":["Recreate one operand with matching dense dims: BCOO((data, indices), shape=...) with the same dense-dimension layout, or use mat.reshape to move dims between sparse/dense","Convert both operands from a canonical source (e.g. fromdense with the same dense_dimensions) so n_dense agrees","If semantically the shapes align, move the mismatched trailing axes so both have the same n_dense count"],"exampleFix":"# before\nlhs = BCOO((d1, i1), shape=(m, n))          # n_dense=0\nrhs = BCOO((d2, i2), shape=(m, n, k))       # n_dense=1\nout = bcoo_multiply_sparse(...)  # NotImplementedError\n# after\nrhs2 = rhs.reshape(m, n * k) or rebuild lhs with n_dense=1 to match:\nlhs = BCOO((d1[..., None], i1), shape=(m, n, 1))  # n_dense=1\nout = lhs * rhs","handlingStrategy":"validation","validationCode":"if lhs.n_dense != rhs.n_dense:\n    rhs = rhs.reshape(rhs.shape[:rhs.n_sparse], rhs.shape[rhs.n_sparse:])  # or rebuild with matching n_dense","typeGuard":"from jax.experimental.sparse import BCOO\ndef compatible_layouts(a: BCOO, b: BCOO) -> bool:\n    return a.n_dense == b.n_dense and a.ndim == b.ndim","tryCatchPattern":"try:\n    out = lhs * rhs\nexcept NotImplementedError as e:\n    if 'dense dimensions' in str(e):\n        rhs = BCOO.fromdense(rhs.todense(), n_batch=rhs.n_batch)  # realign layout\n        out = lhs * rhs\n    else:\n        raise","preventionTips":["Keep a single dense-dimension convention for all BCOO operands in a pipeline","Document n_dense when constructing BCOO tensors; check lhs.n_dense == rhs.n_dense before multiply"],"tags":["jax","sparse","bcoo","not-implemented","dense-dimensions"],"backgroundTag":"unsupported-operand-layout","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}