jax-ml/jax · error · NotImplementedError

sparse rule for {primitive} is not implemented because it wo

Error message

sparse rule for {primitive} is not implemented because it would result in dense output. If this is your intent, use sparse.todense() to convert your arguments to dense matrices.

What it means

When running code under jax.experimental.sparse.sparsify, some LAX primitives (e.g. comparisons like ne, xor, and other _densifying_primitives) applied to sparse operands would produce dense output, which sparsify refuses to do silently.

Source

Thrown at jax/experimental/sparse/transform.py:129

  lax.acos_p,
  lax.acosh_p,
  lax.bessel_i0e_p,
  lax.cos_p,
  lax.cosh_p,
  lax.eq_p,
  lax.exp_p,
  lax.ge_p,
  lax.gt_p,
  lax.le_p,
  lax.lt_p,
  lax.log_p,
  lax.ne_p,
  lax.xor_p
]

def _raise_unimplemented_primitive(primitive):
  if primitive in _densifying_primitives:
    raise NotImplementedError(f"sparse rule for {primitive} is not implemented because it "
                              "would result in dense output. If this is your intent, use "
                              "sparse.todense() to convert your arguments to dense matrices.")
  raise NotImplementedError(f"sparse rule for {primitive} is not implemented.")


Array = Any
ArrayOrSparse = Any


class SparsifyEnv:
  """Environment for sparse jaxpr evaluation.

  The environment is essentially a collection of buffers and/or tracers
  that may be shared between one or more SparsifyValue objects, which
  represent sparse or dense arrays via indices into the list of buffers.
  """
  _buffers : list[Array]

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Convert the sparse argument to dense first with sparse.todense(M) before the op
2. Restructure to use sparsity-aware ops (e.g. operate on M.data instead)
3. Use sparse-aware primitives that preserve sparsity

Example fix

// before
@sparse.sparsify
def f(M):
  return (M != 0).sum()
// after
@sparse.sparsify
def f(M):
  return (sparse.todense(M) != 0).sum()

# inspect supported vs densifying primitives before writing sparsify code
from jax.experimental.sparse import transform
# avoid elementwise comparisons/bitwise ops on sparse values in sparsified fns

When it happens

Trigger: Inside a @sparse.sparsify-decorated function, applying densifying ops (lax.ne, xor, etc.) to a sparse BCOO/BCSR value, e.g. `(M != 0)` or boolean ops on sparse matrices.

Common situations: Comparing sparse matrices elementwise, building masks, or using bitwise ops on sparse values inside sparsified code.

Related errors

• sparse rule for {primitive} is not implemented.
• Sparse metadata format not implemented for {operand_dtype=}
• Unsupported transforms: {a_sparse_metadata_transforms}
• Sparse meta layout loads unsupported.
• Sparse meta layout stores unsupported.

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