jax-ml/jax · error · NotImplementedError

sparse rule for 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.

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 is not implemented because it would result…
• Addition between sparse matrices of different shapes.
• Addition between sparse matrices with different batch/dense…
• axes argument to transpose()
• bcoo_multiply_sparse: arrays with differing numbers of…

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

Appendix: 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)