{"record":{"id":"67c74070d16c3f5b","repo":"jax-ml/jax","slug":"shape-mismatch-sparr-shape-a-shape","errorCode":null,"errorMessage":"shape mismatch: {sparr.shape=} {a.shape=}","messagePattern":"shape mismatch: (.+?) (.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/sparse/bcoo.py","lineNumber":389,"sourceCode":"def bcoo_extract(sparr: BCOO, arr: ArrayLike, *, assume_unique: bool | None = None) -> BCOO:\n  \"\"\"Extract values from a dense array according to the sparse array's indices.\n\n  Args:\n    sparr : BCOO array whose indices will be used for the output.\n    arr : ArrayLike with shape equal to self.shape\n    assume_unique : bool, defaults to sparr.unique_indices\n      If True, extract values for every index, even if index contains duplicates.\n      If False, duplicate indices will have their values summed and returned in\n      the position of the first index.\n\n  Returns:\n    extracted : a BCOO array with the same sparsity pattern as self.\n  \"\"\"\n  if not isinstance(sparr, BCOO):\n    raise TypeError(f\"First argument to bcoo_extract should be a BCOO array. Got {type(sparr)=}\")\n  a = jnp.asarray(arr)\n  if a.shape != sparr.shape:\n    raise ValueError(f\"shape mismatch: {sparr.shape=} {a.shape=}\")\n  if assume_unique is None:\n    assume_unique = sparr.unique_indices\n  data = _bcoo_extract(sparr.indices, a, assume_unique=assume_unique)\n  return BCOO((data, sparr.indices), **sparr._info._asdict())\n\n\ndef _bcoo_extract(indices: Array, arr: Array, *, assume_unique=True) -> Array:\n  \"\"\"Extract BCOO data values from a dense array at given BCOO indices.\n\n  Args:\n    indices: An ndarray; see BCOO indices.\n    arr: A dense array.\n    assume_unique: bool, default=True\n      If True, then indices will be assumed unique and a value will be extracted\n      from arr for each index. Otherwise, extra work will be done to de-duplicate\n      indices to zero-out duplicate extracted values.\n\n  Returns:","sourceCodeStart":371,"sourceCodeEnd":407,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/sparse/bcoo.py#L371-L407","documentation":"bcoo_extract gathers values from a dense array at the sparse array's indices, which requires the dense array to have exactly sparr's shape. It fires when arr.shape differs from the BCOO array's shape, so positions cannot be read.","triggerScenarios":"bcoo_extract(bcoo, dense) where dense.shape != bcoo.shape, e.g. forgetting batch dims, transposed matrix, or leading extra axis.","commonSituations":"Broadcasting assumptions (extract does not broadcast); passing the unbatched dense array to a batched BCOO; shape drift after slicing.","solutions":["Match shapes exactly: reshape/broadcast the dense array first (jnp.broadcast_to(dense, bcoo.shape))","Fix slicing/transposition of the dense operand"],"exampleFix":"// before\nvals = bcoo_extract(bcoo, dense[0])  # forgot batch dim\n// after\nvals = bcoo_extract(bcoo, dense)","handlingStrategy":"validation","validationCode":"if jnp.asarray(arr).shape != sparr.shape:\n    arr = jnp.broadcast_to(arr, sparr.shape)","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Always compare .shape before extract; no broadcasting is applied"],"tags":["jax","sparse","shape-validation","bcoo"],"backgroundTag":"shape-validation-failed","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","contentChangedAt":null,"schemaVersion":2},"datasetVersion":"2026-09-23T08:17:48.524Z"}