jax-ml/jax · error · NotImplementedError

sparse rule for {primitive} is not implemented.

Error message

sparse rule for {primitive} is not implemented.

What it means

A LAX primitive used inside a jax.experimental.sparse.sparsify function has no registered sparse rule at all (it is not in the densifying list either), so sparsify cannot evaluate it on sparse values.

Source

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

  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]

  def __init__(self, bufs=()):
    self._buffers = list(bufs)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Materialize the operand with sparse.todense() before the primitive
2. Check the jax.experimental.sparse.transform.sparse_rules_bcoo/bcsr registries for supported ops
3. Raise a feature request / use dense computation for that portion

Example fix

// before
@sparse.sparsify
def f(M):
  return some_lax_op(M)
// after
@sparse.sparsify
def f(M):
  return some_lax_op(sparse.todense(M))

from jax.experimental.sparse.transform import sparse_rules_bcoo
assert prim in sparse_rules_bcoo, f'{prim} has no sparse rule'

When it happens

Trigger: Calling an exotic lax primitive (no sparse rule registered) on sparse operands inside a @sparse.sparsify function.

Common situations: Using newer/less-common lax ops (e.g. some trig/hyperbolic variants, sort ops) on sparse matrices under sparsify; version drift where a primitive lacks a rule.

Related errors

• sparse rule for {primitive} is not implemented because it wo
• 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/73035ec3d301cc94. Report an issue: GitHub.