{"record":{"id":"244b2cef78abfd79","repo":"jax-ml/jax","slug":"bcsr-sparse-empty-must-have-2-sparse-dimensions","errorCode":null,"errorMessage":"BCSR sparse.empty: must have 2 sparse dimensions.","messagePattern":"BCSR sparse\\.empty: must have 2 sparse dimensions\\.","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/sparse/bcsr.py","lineNumber":927,"sourceCode":"  @classmethod\n  def tree_unflatten(cls, aux_data, children):\n    obj = object.__new__(cls)\n    obj.data, obj.indices, obj.indptr = children\n    if aux_data.keys() != {'shape', 'indices_sorted', 'unique_indices'}:\n      raise ValueError(f\"BCSR.tree_unflatten: invalid {aux_data=}\")\n    obj.__dict__.update(**aux_data)\n    return obj\n\n  @classmethod\n  def _empty(cls, shape, *, dtype=None, index_dtype='int32', n_dense=0,\n             n_batch=0, nse=0):\n    \"\"\"Create an empty BCSR instance. Public method is sparse.empty().\"\"\"\n    shape = tuple(shape)\n    if n_dense < 0 or n_batch < 0 or nse < 0:\n      raise ValueError(f\"Invalid inputs: {shape=}, {n_dense=}, {n_batch=}, {nse=}\")\n    n_sparse = len(shape) - n_dense - n_batch\n    if n_sparse != 2:\n      raise ValueError(\"BCSR sparse.empty: must have 2 sparse dimensions.\")\n    batch_shape, sparse_shape, dense_shape = split_list(shape,\n                                                        [n_batch, n_sparse])\n    data = jnp.zeros((*batch_shape, nse, *dense_shape), dtype)\n    indices = jnp.full((*batch_shape, nse), jnp.array(sparse_shape[1]),\n                       index_dtype)\n    indptr = jnp.zeros((*batch_shape, sparse_shape[0] + 1), index_dtype)\n    return cls((data, indices, indptr), shape=shape)\n\n  def sum_duplicates(self, nse: int | None = None, remove_zeros: bool = True) -> BCSR:\n    \"\"\"Return a copy of the array with duplicate indices summed.\n\n    Additionally, this operation will result in explicit zero entries removed, and\n    indices being sorted in lexicographic order.\n\n    Because the size of the resulting representation depends on the values in the\n    arrays, this operation is not compatible with JIT or other transforms. To use\n    ``sum_duplicates`` in such cases, you may pass a value to `nse` to specify the\n    desired size of the output representation.","sourceCodeStart":909,"sourceCodeEnd":945,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/sparse/bcsr.py#L909-L945","documentation":"jax.experimental.sparse.BCSR (batched CSR) supports exactly 2 sparse dimensions by construction: the CSR format encodes one row-pointer (indptr) dimension and one column-index dimension. sparse.empty()/BCSR._empty computes n_sparse = len(shape) - n_dense - n_batch and rejects anything other than 2. Use BCOO if you need a different number of sparse dimensions.","triggerScenarios":"Calling sparse.empty(shape, format='bcsr') (or BCSR-related empty paths) where len(shape) minus n_dense and n_batch is not 2 — e.g. a 3D shape with n_batch=0, n_dense=0, or a 1D/4D shape, or passing n_dense/n_batch values that leave != 2 sparse dims.","commonSituations":"Migrating code from BCOO (which supports arbitrary n_sparse) to BCSR; passing a batched shape while forgetting to set n_batch so the batch dim is counted as a sparse dim; generically dispatching empty() over many formats with a fixed shape.","solutions":["If leading dims are batch dimensions, pass n_batch equal to the number of batch dims so the remaining sparse part is 2D (e.g. sparse.empty((B, M, N), n_batch=1, format='bcsr'))","If trailing dims are dense, pass n_dense so the sparse part is 2D","If the shape genuinely isn't 2D-sparse (e.g. 3 fully-sparse dims or 1D), use sparse.empty(..., format='bcoo') instead","Catch ValueError and fall back to BCOO when format choice is dynamic"],"exampleFix":"# before\nm = sparse.empty((8, 16, 16), format='bcsr')  # ValueError: must have 2 sparse dimensions\n\n# after (first dim is a batch dim)\nm = sparse.empty((8, 16, 16), n_batch=1, format='bcsr')","handlingStrategy":"validation","validationCode":"from jax.experimental import sparse\nn_sparse = len(shape) - n_dense - n_batch\nassert n_sparse == 2, f'BCSR needs 2 sparse dims, got {n_sparse}; use bcoo'","typeGuard":"def is_bcsr_compatible(shape, n_dense=0, n_batch=0) -> bool:\n    return len(tuple(shape)) - n_dense - n_batch == 2","tryCatchPattern":"try:\n    m = sparse.empty(shape, n_batch=n_batch, format='bcsr')\nexcept ValueError:\n    m = sparse.empty(shape, format='bcoo')","preventionTips":["Compute n_sparse = len(shape) - n_dense - n_batch before format dispatch","Default to bcoo for unknown-dimensionality inputs","Document the 2-sparse-dim invariant wherever BCSR is used"],"tags":["jax","sparse","bcsr","shape-validation"],"backgroundTag":"sparse-format-dimension-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}