jax-ml/jax · error · ValueError

WriteEffect on input buffer {ref_invar_idx}

Error message

WriteEffect on input buffer {ref_invar_idx}

What it means

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.

Source

Thrown at jax/_src/pallas/mosaic/pipeline.py:2287

    if isinstance(avals[ref_idx], state.AbstractRef):
      out_effects.add(ReadEffect(ref_idx) if i < num_inputs else WriteEffect(ref_idx))
  # Propagate effects from `jaxpr`, mapping them to the correct indices in `avals`.
  jaxpr_input_idx = {v: i for i, v in enumerate(
      (*jaxpr.constvars, *jaxpr.invars))}
  for e in jaxpr.effects:
    if not isinstance(e, effects.JaxprInputEffect):
      out_effects.add(e)
      continue
    input_idx = jaxpr_input_idx[e.input]
    if input_idx < len(jaxpr.constvars):
      out_effects.add(e.replace(flat_consts_idx[input_idx]))
    else:
      invar_idx = input_idx - len(jaxpr.constvars)
      if invar_idx < num_ps_leaves:
        continue
      ref_invar_idx = invar_idx - num_ps_leaves
      if ref_invar_idx < num_inputs and isinstance(e, WriteEffect):
        raise ValueError(f"WriteEffect on input buffer {ref_invar_idx}")
      ref_idx = get_ref_idx(flat_refs_idx[ref_invar_idx])
      out_effects.add(e.replace(ref_idx))
  return (), frozenset(out_effects)


# TODO(rdyro): Both primtives require both memory pipeline and core grid
# information which the caching doesn't support yet.
_uncacheable_primitives.add(pipeline_body_p)
_uncacheable_primitives.add(emit_pipeline_p)

def _emit_pipeline_physicalize_rule(
    ctx, *args_flat, body_jaxpr: core.Jaxpr, args_tree, grid_mapping, refs_tree,
    **params
):
  from jax._src.pallas.fuser.fusible_dtype import physicalize_closed_jaxpr  # pyrefly: ignore[missing-import]
  del ctx
  all_args: EmitPipelinePrimitiveArgs = args_tree.unflatten(args_flat)
  with grid_mapping.trace_env():

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Remove writes to input refs in the body and any helper prims
2. Route all outputs through out_specs buffers
3. Inspect the error's ref_invar_idx to find which positional input arg is being written

Example fix

# before
def body(in_ref, out_ref):
  in_ref[...] = f(in_ref[...])
# after
def body(in_ref, out_ref):
  out_ref[...] = f(in_ref[...])

When it happens

Trigger: 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.

Common situations: Identical user mistakes as 2236: mutating in_refs, sharing refs between inputs and outputs, or pallas primitives with write effects targeting an input.

Related errors

• WriteEffect should not apply to an input buffer {ref_invar_i
• Effects not supported in `custom_jvp`: {disallowed_effects}
• Effects not supported in `custom_jvp`: {disallowed}
• Cannot lower jaxpr with effects: {closed_jaxpr.effects}
• `JaxprInputEffect` {eff} does not have a corresponding jaxpr

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