jax-ml/jax · error · ValueError

Cannot transpose with respect to sparse indices

Error message

Cannot transpose with respect to sparse indices

What it means

During JAX's transpose (backward-pass) rule for csr_todense, the CSR index arrays (indices, indptr) themselves were traced as differentiable values. Gradients with respect to sparse structure are not defined/differentiable, so JAX raises.

Source

Thrown at jax/experimental/sparse/csr.py:288

  dtype = data_aval.dtype
  if not (np.issubdtype(dtype, np.floating) or np.issubdtype(dtype, np.complexfloating)):
    warnings.warn(f"csr_todense cusparse/hipsparse lowering not available for {dtype=}. "
                  "Falling back to default implementation.", CuSparseEfficiencyWarning)
    return _csr_todense_lowering(ctx, data, indices, indptr, shape=shape)
  return [_lowerings.csr_todense_gpu_lowering(
      ctx, data, indices, indptr, shape=shape,
      target_name_prefix=target_name_prefix)]


def _csr_todense_jvp(data_dot, data, indices, indptr, *, shape):
  return _csr_todense(data_dot, indices, indptr, shape=shape)

def _csr_todense_transpose(ct, data, indices, indptr, *, shape):
  # Note: we assume that transpose has the same sparsity pattern.
  # Can we check this?
  assert ad.is_undefined_primal(data)
  if ad.is_undefined_primal(indices) or ad.is_undefined_primal(indptr):
    raise ValueError("Cannot transpose with respect to sparse indices")
  assert ct.shape == shape
  assert indices.aval.dtype == indptr.aval.dtype
  assert ct.dtype == data.aval.dtype
  return _csr_extract(indices, indptr, ct), indices, indptr

ad.defjvp(csr_todense_p, _csr_todense_jvp, None, None)
ad.primitive_transposes[csr_todense_p] = _csr_todense_transpose
mlir.register_lowering(csr_todense_p, _csr_todense_lowering)
dispatch.simple_impl(csr_todense_p)

mlir.register_lowering(
    csr_todense_p,
    partial(_csr_todense_gpu_lowering, target_name_prefix='cu'),
    platform='cuda')
mlir.register_lowering(
    csr_todense_p,
    partial(_csr_todense_gpu_lowering, target_name_prefix='hip'),
    platform='rocm')

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use jax.lax.stop_gradient on indices/indptr before passing them into the sparse op
2. Keep index arrays as static/non-traced inputs (construct CSR outside grad scope)
3. Differentiate only with respect to the data buffer

Example fix

// before
M = CSR((data, idx, indptr))
out = jax.grad(lambda d: todense(CSR((d, idx, indptr))).sum())(data)  # idx traced
// after
idx = jax.lax.stop_gradient(idx)
indptr = jax.lax.stop_gradient(indptr)
out = jax.grad(lambda d: todense(CSR((d, idx, indptr))).sum())(data)

import jax
assert not isinstance(indices, jax.core.Tracer), 'indices must not be traced'
assert not isinstance(indptr, jax.core.Tracer), 'indptr must not be traced'

When it happens

Trigger: Calling jax.grad (or vjp/jvp backward) on a function whose csr_todense (e.g. sparse.CSR.dot, todense) has indices/indptr as traced/differentiated arguments, e.g. differentiating through code that produces index arrays from parameters.

Common situations: Wrapping sparse construction (jnp.argsort, searchsorted outputs feeding indices) inside a differentiated function; using lax.custom_sparse operations where indices depend on inputs.

Related errors

• todense_transpose for {type(obj)}
• Buffer callbacks do not support transpose. Please use `jax.c
• Called {name} with a float0 array. float0s do not support an
• transpose of bcoo_extract with assume_unique=False
• CSR must have ndim=2; got {shape=}

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