{"record":{"id":"79c11d897b5f525c","repo":"jax-ml/jax","slug":"writeeffect-on-input-buffer-ref-invar-idx","errorCode":null,"errorMessage":"WriteEffect on input buffer {ref_invar_idx}","messagePattern":"WriteEffect on input buffer (.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/pipeline.py","lineNumber":2287,"sourceCode":"    if isinstance(avals[ref_idx], state.AbstractRef):\n      out_effects.add(ReadEffect(ref_idx) if i < num_inputs else WriteEffect(ref_idx))\n  # Propagate effects from `jaxpr`, mapping them to the correct indices in `avals`.\n  jaxpr_input_idx = {v: i for i, v in enumerate(\n      (*jaxpr.constvars, *jaxpr.invars))}\n  for e in jaxpr.effects:\n    if not isinstance(e, effects.JaxprInputEffect):\n      out_effects.add(e)\n      continue\n    input_idx = jaxpr_input_idx[e.input]\n    if input_idx < len(jaxpr.constvars):\n      out_effects.add(e.replace(flat_consts_idx[input_idx]))\n    else:\n      invar_idx = input_idx - len(jaxpr.constvars)\n      if invar_idx < num_ps_leaves:\n        continue\n      ref_invar_idx = invar_idx - num_ps_leaves\n      if ref_invar_idx < num_inputs and isinstance(e, WriteEffect):\n        raise ValueError(f\"WriteEffect on input buffer {ref_invar_idx}\")\n      ref_idx = get_ref_idx(flat_refs_idx[ref_invar_idx])\n      out_effects.add(e.replace(ref_idx))\n  return (), frozenset(out_effects)\n\n\n# TODO(rdyro): Both primtives require both memory pipeline and core grid\n# information which the caching doesn't support yet.\n_uncacheable_primitives.add(pipeline_body_p)\n_uncacheable_primitives.add(emit_pipeline_p)\n\ndef _emit_pipeline_physicalize_rule(\n    ctx, *args_flat, body_jaxpr: core.Jaxpr, args_tree, grid_mapping, refs_tree,\n    **params\n):\n  from jax._src.pallas.fuser.fusible_dtype import physicalize_closed_jaxpr  # pyrefly: ignore[missing-import]\n  del ctx\n  all_args: EmitPipelinePrimitiveArgs = args_tree.unflatten(args_flat)\n  with grid_mapping.trace_env():","sourceCodeStart":2269,"sourceCodeEnd":2305,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/pipeline.py#L2269-L2305","documentation":"Same rule as its sibling check in a different jaxpr-rewriting path: a WriteEffect on an input buffer is rejected. Here it surfaces during the grid/iteration jaxpr rewriting stage rather than body effect inference, with a shorter message naming the offending input index.","triggerScenarios":"An input buffer of the pipeline receives a write somewhere in the traced iteration logic — typically induced by the same body writing to an in_ref, but caught in the other rewriting pass.","commonSituations":"Identical user mistakes as 2236: mutating in_refs, sharing refs between inputs and outputs, or pallas primitives with write effects targeting an input.","solutions":["Remove writes to input refs in the body and any helper prims","Route all outputs through out_specs buffers","Inspect the error's ref_invar_idx to find which positional input arg is being written"],"exampleFix":"# before\ndef body(in_ref, out_ref):\n  in_ref[...] = f(in_ref[...])\n# after\ndef body(in_ref, out_ref):\n  out_ref[...] = f(in_ref[...])","handlingStrategy":"type-guard","validationCode":null,"typeGuard":null,"tryCatchPattern":null,"preventionTips":["Keep inputs read-only; use scratch/out buffers for mutation","Use ref_invar_idx from the message to identify which positional input is written"],"tags":["jax","pallas","effects","input-mutation"],"backgroundTag":"write-to-read-only","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}