{"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 requires the dense array's shape to equal the BCOO's shape, since it extracts values at matching indices. Any mismatch raises ValueError showing both shapes.","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","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}