jax-ml/jax · error · NotImplementedError

transpose of bcoo_extract with assume_unique=False

Error message

transpose of bcoo_extract with assume_unique=False

What it means

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.

Source

Thrown at jax/experimental/sparse/bcoo.py:466

    for _ in range(props.n_batch):
      f = vmap(f)
    result = f(result, sort_ind)
  return result

@bcoo_extract_p.def_abstract_eval
def _bcoo_extract_abstract_eval(indices, arr, *, assume_unique):
  _ = bool(assume_unique)
  n_batch, _, n_dense, nse = _validate_bcoo_indices(indices, arr.shape)
  out_shape = arr.shape[:n_batch] + (nse,) + arr.shape[arr.ndim - n_dense:]
  return core.ShapedArray(out_shape, arr.dtype)

def _bcoo_extract_jvp(arr_dot, indices, arr, *, assume_unique):
  assert arr_dot.shape == arr.shape
  return _bcoo_extract(indices, arr_dot, assume_unique=assume_unique)

def _bcoo_extract_transpose(ct, indices, arr, *, assume_unique):
  if not assume_unique:
    raise NotImplementedError("transpose of bcoo_extract with assume_unique=False")
  assert ad.is_undefined_primal(arr)
  if ad.is_undefined_primal(indices):
    raise ValueError("Cannot transpose with respect to sparse indices")
  assert ct.dtype == arr.aval.dtype
  return indices, _bcoo_todense(ct, indices, spinfo=SparseInfo(arr.aval.shape))

def _bcoo_extract_batching_rule(batched_args, batch_dims, *, assume_unique):
  indices, arr = batched_args
  assert any(b is not None for b in batch_dims)
  if batch_dims[0] is None:
    bdim = batch_dims[1]
    indices = lax.expand_dims(indices, (bdim,))
  elif batch_dims[1] is None:
    # TODO(jakevdp) can we handle this case without explicit broadcasting?
    bdim = batch_dims[0]
    result_shape = list(arr.shape)
    result_shape.insert(bdim, indices.shape[bdim])
    arr = lax.broadcast_in_dim(arr, result_shape, (bdim,))

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Call bcoo_sum_duplicates on the BCOO first so unique_indices becomes True
2. Pass assume_unique=True explicitly if you know indices are unique

Example fix

// before
vals = bcoo_extract(bcoo, dense)
jax.grad(f)(...)  # bcoo has duplicate indices
// after
bcoo = bcoo.sum_duplicates()
vals = bcoo_extract(bcoo, dense)
jax.grad(f)(...)

if not bcoo.unique_indices:
    bcoo = bcoo.sum_duplicates()

try:
    jax.grad(f)(x)
except NotImplementedError:
    sp = sp.sum_duplicates(); jax.grad(f)(x)

When it happens

Trigger: 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.

Common situations: BCOO built by concatenation or manual buffers where duplicates exist; grad of sampled matmul (spdot with a dense mask) using non-unique indices.

Related errors

• transpose permutation must be a tuple/list/ndarray, got {typ
• transpose permutation isn't a permutation of operand dimensi
• transpose permutation cannot permute batch axes with non-bat
• Pure callbacks do not support transpose. Please use `jax.cus
• for transpose support, subclass {type(self)} must implement

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/ad7a19c04da92a4b. Report an issue: GitHub.