{"record":{"id":"ff299ade2dafd953","repo":"jax-ml/jax","slug":"expected-0-bdim-n-batch-got-bdim-n-batc","errorCode":null,"errorMessage":"Expected 0 < bdim <= n_batch; got {bdim=}, {n_batch=}","messagePattern":"Expected 0 < bdim <= n_batch; got (.+?), (.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/sparse/bcoo.py","lineNumber":356,"sourceCode":"\n  return primals_out, tangents_out\n\ndef _bcoo_fromdense_transpose(ct, M, *, nse, n_batch, n_dense, index_dtype):\n  data, indices = ct\n  n_sparse = M.ndim - n_batch - n_dense\n  assert data.shape == M.shape[:n_batch] + (nse,) + M.shape[n_batch + n_sparse:]\n  assert indices.shape == M.shape[:n_batch] + (n_sparse, nse)\n  assert indices.dtype == index_dtype\n  if isinstance(indices, ad.Zero):\n    raise ValueError(\"Cannot transpose with respect to sparse indices\")\n  assert ad.is_undefined_primal(M)\n  return _bcoo_todense(data, indices, spinfo=SparseInfo(M.aval.shape))\n\ndef _bcoo_fromdense_batching_rule(batched_args, batch_dims, *, nse, n_batch, n_dense, index_dtype):\n  M, = batched_args\n  bdim, = batch_dims\n  if not (0 <= bdim <= n_batch):\n    raise ValueError(f\"Expected 0 < bdim <= n_batch; got {bdim=}, {n_batch=}\")\n  return _bcoo_fromdense(M, nse=nse, n_batch=n_batch + 1, n_dense=n_dense, index_dtype=index_dtype), (bdim, bdim)\n\nad.primitive_jvps[bcoo_fromdense_p] = _bcoo_fromdense_jvp\nad.primitive_transposes[bcoo_fromdense_p] = _bcoo_fromdense_transpose\nbatching.primitive_batchers[bcoo_fromdense_p] = _bcoo_fromdense_batching_rule\nmlir.register_lowering(bcoo_fromdense_p, mlir.lower_fun(\n    _bcoo_fromdense_impl, multiple_results=True))\n\n#----------------------------------------------------------------------\n# bcoo_extract\n\nbcoo_extract_p = core.Primitive('bcoo_extract')\n\n\ndef 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:","sourceCodeStart":338,"sourceCodeEnd":374,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/sparse/bcoo.py#L338-L374","documentation":"When vmap-ing bcoo_fromdense, the batched axis of the dense input must be within [0, n_batch]. The code even has an off-by-one message bug (checks 0 <= bdim <= n_batch but the message says 0 < bdim). Raised for out-of-range batch dims.","triggerScenarios":"jax.vmap over an axis beyond the declared n_batch when converting dense to BCOO, e.g. vmap-ing a fromdense call where the mapped dim falls in sparse dimensions.","commonSituations":"vmap over rows of a matrix being converted elementwise to sparse; converting dense arrays whose leading dims aren't declared as batch dims.","solutions":["Declare the mapped dimension as a batch dim: pass n_batch including it in BCOO.fromdense / bcoo_fromdense","Move the mapped axis to position 0 and set n_batch accordingly","Do the dense-to-sparse conversion outside vmap, then vmap over the resulting BCOO"],"exampleFix":"// before\nf = jax.vmap(lambda x: BCOO.fromdense(x, nse=4))(batched_dense)  # n_batch mismatch\n// after\nf = BCOO.fromdense(batched_dense, nse=4, n_batch=1)  # batch conversion directly","handlingStrategy":"validation","validationCode":"assert 0 <= mapped_axis <= n_batch_declared, 'bdim outside batch dims'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Declare the mapped axis in n_batch of fromdense","Convert batched dense arrays with a single fromdense call instead of vmap"],"tags":["jax","sparse","bcoo","vmap","fromdense"],"backgroundTag":"vmap-batch-dim-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}