{"record":{"id":"f9b99033a943964e","repo":"jax-ml/jax","slug":"writeeffect-should-not-apply-to-an-input-buffer-r","errorCode":null,"errorMessage":"WriteEffect should not apply to an input buffer {ref_invar_idx} in pipeline body jaxpr: {body_jaxpr}","messagePattern":"WriteEffect should not apply to an input buffer (.+?) in pipeline body jaxpr: (.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/pipeline.py","lineNumber":2229,"sourceCode":"  # Propagate effects from `body_jaxpr`, mapping them to the correct indices in\n  # `avals`.\n  body_input_idx = {v: i for i, v in enumerate(\n      (*body_jaxpr.constvars, *body_jaxpr.invars))}\n  for e in body_jaxpr.effects:\n    if not isinstance(e, effects.JaxprInputEffect):\n      out_effects.add(e)\n      continue\n    input_idx = body_input_idx[e.input]\n    if input_idx < len(body_jaxpr.constvars):\n      const_offset = all_args.body_offset\n      out_effects.add(e.replace(const_offset + input_idx))\n    else:\n      invar_idx = input_idx - len(body_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(\n            f\"WriteEffect should not apply to an input buffer {ref_invar_idx} in\"\n            f\" pipeline body jaxpr: {body_jaxpr}\")\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# TODO(rdyro): Either generalize or merge with another primitive. This primitive\n# perfoms an \"eval jaxpr\" operation, but is currently tailored to calling the\n# pipeline body in the emit_pipeline primtiive - it resolves TransformedRefs and\n# binds the user grid indices to lowering.\n# This primitive is specialized to resolve TransformedRefs passed as arguments\n# and evaluate the body jaxpr with the resolved Refs because it assumes the body\n# was traced \"generically\" with Refs. However, the emit_pipeline is allowed to\n# pass in TransformedRefs as arguments to the body.\npipeline_body_p = core.Primitive(\"pipeline_body\")\npipeline_body_p.multiple_results = True\n\n@pipeline_body_p.def_effectful_abstract_eval","sourceCodeStart":2211,"sourceCodeEnd":2247,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/pipeline.py#L2211-L2247","documentation":"When JAX infers effects for the pipelined body, an input buffer (a read-only in_ref) must not carry a WriteEffect. If the effect-remapping logic finds a WriteEffect attached to an input ref, the body wrote to an input, which the pipeline forbids; the message includes the jaxpr for debugging.","triggerScenarios":"Writing to an input ref inside the kernel body, e.g. `in_ref[...] = ...` or passing an input ref to a helper that mutates it, in a lower-level jaxpr-based pipeline path.","commonSituations":"Porting a kernel that used a single scratch ref for both input and output; in-place normalization or caching patterns applied to input buffers.","solutions":["Never assign to in_refs; allocate an out_spec buffer (or scratch) and write results there","If in-place semantics are needed, declare that buffer as an output spec with the matching block mapping","Check helper functions (e.g. custom update routines) for hidden writes to input refs"],"exampleFix":"# before\ndef body(in_ref, out_ref, i):\n  in_ref[i] *= scale  # WRONG: writes input\n  out_ref[i] = in_ref[i]\n# after\ndef body(in_ref, out_ref, i):\n  out_ref[i] = in_ref[i] * scale","handlingStrategy":"type-guard","validationCode":"import jax\nres = jax.eval_shape(lambda: None)  # ensure no writes: inspect body jaxpr effects in tests\neffects = getattr(traced_body, 'effects', None)\nassert not any(isinstance(e, jax.experimental.pallas.effects.WriteEffect) for e in (effects or []))","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Never assign to in_refs in kernel bodies","Route results exclusively through out_spec buffers"],"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"}