{"record":{"id":"ad7a19c04da92a4b","repo":"jax-ml/jax","slug":"transpose-of-bcoo-extract-with-assume-unique-false","errorCode":null,"errorMessage":"transpose of bcoo_extract with assume_unique=False","messagePattern":"transpose of bcoo_extract with assume_unique=False","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/experimental/sparse/bcoo.py","lineNumber":466,"sourceCode":"    for _ in range(props.n_batch):\n      f = vmap(f)\n    result = f(result, sort_ind)\n  return result\n\n@bcoo_extract_p.def_abstract_eval\ndef _bcoo_extract_abstract_eval(indices, arr, *, assume_unique):\n  _ = bool(assume_unique)\n  n_batch, _, n_dense, nse = _validate_bcoo_indices(indices, arr.shape)\n  out_shape = arr.shape[:n_batch] + (nse,) + arr.shape[arr.ndim - n_dense:]\n  return core.ShapedArray(out_shape, arr.dtype)\n\ndef _bcoo_extract_jvp(arr_dot, indices, arr, *, assume_unique):\n  assert arr_dot.shape == arr.shape\n  return _bcoo_extract(indices, arr_dot, assume_unique=assume_unique)\n\ndef _bcoo_extract_transpose(ct, indices, arr, *, assume_unique):\n  if not assume_unique:\n    raise NotImplementedError(\"transpose of bcoo_extract with assume_unique=False\")\n  assert ad.is_undefined_primal(arr)\n  if ad.is_undefined_primal(indices):\n    raise ValueError(\"Cannot transpose with respect to sparse indices\")\n  assert ct.dtype == arr.aval.dtype\n  return indices, _bcoo_todense(ct, indices, spinfo=SparseInfo(arr.aval.shape))\n\ndef _bcoo_extract_batching_rule(batched_args, batch_dims, *, assume_unique):\n  indices, arr = batched_args\n  assert any(b is not None for b in batch_dims)\n  if batch_dims[0] is None:\n    bdim = batch_dims[1]\n    indices = lax.expand_dims(indices, (bdim,))\n  elif batch_dims[1] is None:\n    # TODO(jakevdp) can we handle this case without explicit broadcasting?\n    bdim = batch_dims[0]\n    result_shape = list(arr.shape)\n    result_shape.insert(bdim, indices.shape[bdim])\n    arr = lax.broadcast_in_dim(arr, result_shape, (bdim,))","sourceCodeStart":448,"sourceCodeEnd":484,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/sparse/bcoo.py#L448-L484","documentation":"The transpose (backward-pass) of bcoo_extract is only defined when indices are unique; with assume_unique=False a correct scatter-based transpose is not implemented, so it raises NotImplementedError.","triggerScenarios":"Differentiating through bcoo_extract (directly or via bcoo_dot_general_sampled) on a BCOO whose unique_indices is False and assume_unique not explicitly set True.","commonSituations":"BCOO built by concatenation or manual buffers where duplicates exist; grad of sampled matmul (spdot with a dense mask) using non-unique indices.","solutions":["Call bcoo_sum_duplicates on the BCOO first so unique_indices becomes True","Pass assume_unique=True explicitly if you know indices are unique"],"exampleFix":"// before\nvals = bcoo_extract(bcoo, dense)\njax.grad(f)(...)  # bcoo has duplicate indices\n// after\nbcoo = bcoo.sum_duplicates()\nvals = bcoo_extract(bcoo, dense)\njax.grad(f)(...)","handlingStrategy":"validation","validationCode":"if not bcoo.unique_indices:\n    bcoo = bcoo.sum_duplicates()","typeGuard":null,"tryCatchPattern":"try:\n    jax.grad(f)(x)\nexcept NotImplementedError:\n    sp = sp.sum_duplicates(); jax.grad(f)(x)","preventionTips":["Call sum_duplicates before differentiating through extract","Set assume_unique=True only when indices are truly unique"],"tags":["jax","sparse","bcoo","autodiff","transpose"],"backgroundTag":"unsupported-autodiff-rule","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}