{"record":{"id":"5d7b4f7cd1fd6fe0","repo":"jax-ml/jax","slug":"bcoo-extract-with-unequal-batch-dimensions","errorCode":null,"errorMessage":"bcoo_extract with unequal batch dimensions.","messagePattern":"bcoo_extract with unequal batch dimensions\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/experimental/sparse/bcoo.py","lineNumber":487,"sourceCode":"    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,))\n  else:\n    if batch_dims[0] != batch_dims[1]:\n      raise NotImplementedError(\"bcoo_extract with unequal batch dimensions.\")\n    bdim = batch_dims[0]\n  n_batch = indices.ndim - 2\n  if bdim >= n_batch:\n    raise ValueError(f\"{batch_dims=} out of range for indices with {n_batch=}\")\n  return _bcoo_extract(indices, arr, assume_unique=assume_unique), bdim\n\nad.defjvp(bcoo_extract_p, None, _bcoo_extract_jvp)\nad.primitive_transposes[bcoo_extract_p] = _bcoo_extract_transpose\nbatching.primitive_batchers[bcoo_extract_p] = _bcoo_extract_batching_rule\nmlir.register_lowering(bcoo_extract_p, mlir.lower_fun(\n    _bcoo_extract_impl, multiple_results=False))\n\n#----------------------------------------------------------------------\n# bcoo_transpose\n# transpose of a BCOO array\n\nbcoo_transpose_p = core.Primitive('bcoo_transpose')\nbcoo_transpose_p.multiple_results = True","sourceCodeStart":469,"sourceCodeEnd":505,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/sparse/bcoo.py#L469-L505","documentation":"When vmap-ing bcoo_extract with both indices and dense array batched, their batch dims must be equal; unequal batch dims (e.g. one batched at axis 0, the other at axis 1) are not implemented and raise NotImplementedError.","triggerScenarios":"jax.vmap(bcoo_extract, in_axes=(0, 1)) or any vmap configuration yielding different batch_dims for indices and arr inside the batching rule.","commonSituations":"Batching the sparse operand along one axis and the dense operand along another; in_axes mismatch when mapping lists of matrices.","solutions":["Use matching in_axes for both operands, e.g. in_axes=0","Move both operands' mapped axes to the same position with jnp.moveaxis before vmap"],"exampleFix":"// before\nout = jax.vmap(bcoo_extract, in_axes=(0, 1))(bcoo_batch, dense_batch)\n// after\ndense_batch = jnp.moveaxis(dense_batch, 1, 0)\nout = jax.vmap(bcoo_extract, in_axes=0)(bcoo_batch, dense_batch)","handlingStrategy":"validation","validationCode":"assert in_axes[0] == in_axes[1] or None in in_axes, 'batch dims must match'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Use identical in_axes for indices and dense operands in vmap of extract"],"tags":["jax","sparse","bcoo","vmap","batching"],"backgroundTag":"vmap-batch-dim-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}