{"record":{"id":"31140b4ecb263a12","repo":"jax-ml/jax","slug":"sparse-rule-for-lax-integer-pow-p-with-non-posit","errorCode":null,"errorMessage":"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.","messagePattern":"sparse rule for (.+?) with non-positive exponent (.+?) 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\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/experimental/sparse/transform.py","lineNumber":570,"sourceCode":"  lax.slice_p: sparse.bcoo_slice,\n  lax.squeeze_p: sparse.bcoo_squeeze,\n}\n\nfor prim, bcoo_impl in _BCOO_STANDARD_PRIMITIVES.items():\n  sparse_rules_bcoo[prim] = _standard_sparse_rule(prim, bcoo_impl)\n\n_BCSR_STANDARD_PRIMITIVES = {\n  lax.dot_general_p: sparse.bcsr_dot_general,\n  lax.broadcast_in_dim_p: sparse.bcsr_broadcast_in_dim,\n  lax.concatenate_p: lambda *a, **k: sparse.bcsr_concatenate(a, **k),\n}\n\nfor prim, bcsr_impl in _BCSR_STANDARD_PRIMITIVES.items():\n  sparse_rules_bcsr[prim] = _standard_sparse_rule(prim, bcsr_impl)\n\ndef _integer_pow_sparse(spenv, *spvalues, y):\n  if y <= 0:\n    raise NotImplementedError(f\"sparse rule for {lax.integer_pow_p} with non-positive exponent {y} is \"\n                              \"not implemented because it would result in dense output. If this is your \"\n                              \"intent, use sparse.todense() to convert your argument to a dense array.\")\n  return _zero_preserving_unary_op(lax.integer_pow_p, False)(spenv, *spvalues, y=y)\n\nsparse_rules_bcoo[lax.integer_pow_p] = _integer_pow_sparse\nsparse_rules_bcsr[lax.integer_pow_p] = _integer_pow_sparse\n\ndef _transpose_sparse(spenv, *spvalues, permutation):\n  permutation = tuple(permutation)\n  args = spvalues_to_arrays(spenv, spvalues)\n  shape = args[0].shape\n  mat_transposed = sparse.bcoo_transpose(args[0], permutation=permutation)\n  out_shape = tuple(shape[i] for i in permutation)\n\n  n_batch = args[0].indices.ndim - 2\n  n_sparse = args[0].indices.shape[-1]\n  batch_dims_unchanged = (permutation[:n_batch] == tuple(range(n_batch)))\n  dense_dims_unchanged = (permutation[n_batch + n_sparse:] == tuple(range(n_batch + n_sparse, len(shape))))","sourceCodeStart":552,"sourceCodeEnd":588,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/sparse/transform.py#L552-L588","documentation":"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.","triggerScenarios":"Computing M ** 0, M ** -1, M ** -2, or 1 / M equivalently on a sparse value inside a @sparse.sparsify function.","commonSituations":"Normalizing a sparse matrix by its inverse or power; writing (1/M) style regularization inside sparsified code.","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"],"exampleFix":"// before\n@sparse.sparsify\ndef f(M):\n  return M ** -1\n// after\n@sparse.sparsify\ndef f(M):\n  return 1.0 / sparse.todense(M)","handlingStrategy":"validation","validationCode":"assert y > 0, 'integer powers with y <= 0 densify sparse values; use todense first'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Avoid **0 / **negative on sparse values inside sparsify","Prefer operating on M.data for reciprocal-style transforms"],"tags":["jax","sparse","sparsify","power","densifying"],"backgroundTag":"sparse-autodiff-unsupported-op","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}