{"record":{"id":"a4bdee580238bc03","repo":"jax-ml/jax","slug":"invalid-data-shape-for-nse-n-batch-n-d","errorCode":null,"errorMessage":"Invalid {data.shape=} for {nse=}, {n_batch=}, {n_dense=}","messagePattern":"Invalid (.+?) for (.+?), (.+?), (.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/sparse/bcoo.py","lineNumber":140,"sourceCode":"  n_sparse: int\n  n_dense: int\n  nse: int\n\nclass Buffer(Protocol):\n  @property\n  def shape(self) -> Shape: ...\n  @property\n  def dtype(self) -> Any: ...\n\n\ndef _validate_bcoo(data: Buffer, indices: Buffer, shape: Sequence[int]) -> BCOOProperties:\n  props = _validate_bcoo_indices(indices, shape)\n  n_batch, n_sparse, n_dense, nse = props\n  shape = tuple(shape)\n  if any(s1 not in (1, s2) for s1, s2 in safe_zip(data.shape[:n_batch], shape[:n_batch])):\n    raise ValueError(f\"data batch dimensions not compatible for {data.shape=}, {shape=}\")\n  if data.shape[n_batch:] != (nse,) + shape[n_batch + n_sparse:]:\n    raise ValueError(f\"Invalid {data.shape=} for {nse=}, {n_batch=}, {n_dense=}\")\n  return props\n\n\ndef _validate_bcoo_indices(indices: Buffer, shape: Sequence[int]) -> BCOOProperties:\n  assert jnp.issubdtype(indices.dtype, jnp.integer)\n  shape = tuple(shape)\n  nse, n_sparse = indices.shape[-2:]\n  n_batch = len(indices.shape) - 2\n  n_dense = len(shape) - n_batch - n_sparse\n  assert n_dense >= 0\n  if any(s1 not in (1, s2) for s1, s2 in safe_zip(indices.shape[:n_batch], shape[:n_batch])):\n    raise ValueError(f\"indices batch dimensions not compatible for {indices.shape=}, {shape=}\")\n  if indices.shape[n_batch:] != (nse, n_sparse):\n    raise ValueError(f\"Invalid ={indices.shape=} for {nse=}, {n_batch=}, {n_dense=}\")\n  return BCOOProperties(n_batch=n_batch, n_sparse=n_sparse, n_dense=n_dense, nse=nse)\n\n\n#----------------------------------------------------------------------","sourceCodeStart":122,"sourceCodeEnd":158,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/sparse/bcoo.py#L122-L158","documentation":"Beyond the batch dims, a BCOO's data must have shape (nse,) + dense dims, where nse comes from indices.shape[-2]. If it doesn't, _validate_bcoo raises this ValueError.","triggerScenarios":"Building BCOO where data.shape[n_batch:] != (nse,) + shape[n_batch+n_sparse:], e.g. data has length 10 but indices imply nse=8, or dense dims missing.","commonSituations":"Trimming data without trimming indices; wrong nse after eliminate_zeros/sum_duplicates done manually; forgetting trailing dense dims for dense-sparse hybrids.","solutions":["Make data length equal indices.shape[-2] (nse)","Regenerate the pair via BCOO.fromdense(...) or jax.experimental.sparse.bcoo.bcoo_todense round-trip","Use bcoo_update_layout to fix layout instead of hand-editing buffers"],"exampleFix":"// before\nm = BCOO((data[:5], indices), shape=shape)  # indices has nse=8\n// after\nm = BCOO((data[:5], indices[:, :5]), shape=shape)","handlingStrategy":"validation","validationCode":"nse = indices.shape[-2]\nassert data.shape[n_batch:] == (nse,) + shape[n_batch+n_sparse:], 'data/nse mismatch'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Never trim data without trimming indices","Use bcoo_sum_duplicates to keep buffers consistent"],"tags":["jax","sparse","bcoo","shape-validation","nse"],"backgroundTag":"shape-validation-failed","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}