jax-ml/jax · error · NotImplementedError

sparse rule for {lax.integer_pow_p} with non-positive expone

Q: What causes "sparse rule for {lax.integer_pow_p} with non-positive exponent {y} is not implemented because it would result in dense output. If this is your intent, use sparse.todense() to convert your argument to a dense array." in jax-ml/jax?

Use a positive exponent, or convert with sparse.todense(M) before exponentiating

Error message

sparse rule for {lax.integer_pow_p} with non-positive exponent {y} is not implemented because it would result in dense output. If this is your intent, use sparse.todense() to convert your argument to a dense array.

What it means

Under sparsify, x ** y with lax.integer_pow_p and a non-positive exponent (y <= 0) would densify a sparse matrix (zeros become infinities/ones), so it is deliberately not implemented for sparse operands.

Source

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

  lax.slice_p: sparse.bcoo_slice,
  lax.squeeze_p: sparse.bcoo_squeeze,
}

for prim, bcoo_impl in _BCOO_STANDARD_PRIMITIVES.items():
  sparse_rules_bcoo[prim] = _standard_sparse_rule(prim, bcoo_impl)

_BCSR_STANDARD_PRIMITIVES = {
  lax.dot_general_p: sparse.bcsr_dot_general,
  lax.broadcast_in_dim_p: sparse.bcsr_broadcast_in_dim,
  lax.concatenate_p: lambda *a, **k: sparse.bcsr_concatenate(a, **k),
}

for prim, bcsr_impl in _BCSR_STANDARD_PRIMITIVES.items():
  sparse_rules_bcsr[prim] = _standard_sparse_rule(prim, bcsr_impl)

def _integer_pow_sparse(spenv, *spvalues, y):
  if y <= 0:
    raise NotImplementedError(f"sparse rule for {lax.integer_pow_p} with non-positive exponent {y} is "
                              "not implemented because it would result in dense output. If this is your "
                              "intent, use sparse.todense() to convert your argument to a dense array.")
  return _zero_preserving_unary_op(lax.integer_pow_p, False)(spenv, *spvalues, y=y)

sparse_rules_bcoo[lax.integer_pow_p] = _integer_pow_sparse
sparse_rules_bcsr[lax.integer_pow_p] = _integer_pow_sparse

def _transpose_sparse(spenv, *spvalues, permutation):
  permutation = tuple(permutation)
  args = spvalues_to_arrays(spenv, spvalues)
  shape = args[0].shape
  mat_transposed = sparse.bcoo_transpose(args[0], permutation=permutation)
  out_shape = tuple(shape[i] for i in permutation)

  n_batch = args[0].indices.ndim - 2
  n_sparse = args[0].indices.shape[-1]
  batch_dims_unchanged = (permutation[:n_batch] == tuple(range(n_batch)))
  dense_dims_unchanged = (permutation[n_batch + n_sparse:] == tuple(range(n_batch + n_sparse, len(shape))))

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Use a positive exponent, or convert with sparse.todense(M) before exponentiating
Express the operation on M.data only if the sparsity semantics allow
Add a small epsilon / restructure so the exponent is positive

Example fix

// before
@sparse.sparsify
def f(M):
  return M ** -1
// after
@sparse.sparsify
def f(M):
  return 1.0 / sparse.todense(M)

Defensive patterns

Strategy: validation

Validate before calling

assert y > 0, 'integer powers with y <= 0 densify sparse values; use todense first'

Prevention

Avoid **0 / **negative on sparse values inside sparsify
Prefer operating on M.data for reciprocal-style transforms

When it happens

Trigger: Computing M ** 0, M ** -1, M ** -2, or 1 / M equivalently on a sparse value inside a @sparse.sparsify function.

Common situations: Normalizing a sparse matrix by its inverse or power; writing (1/M) style regularization inside sparsified code.

Related errors

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