jax-ml/jax · error · ValueError

WriteEffect should not apply to an input buffer {ref_invar_i

Error message

WriteEffect should not apply to an input buffer {ref_invar_idx} in pipeline body jaxpr: {body_jaxpr}

What it means

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.

Source

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

  # Propagate effects from `body_jaxpr`, mapping them to the correct indices in
  # `avals`.
  body_input_idx = {v: i for i, v in enumerate(
      (*body_jaxpr.constvars, *body_jaxpr.invars))}
  for e in body_jaxpr.effects:
    if not isinstance(e, effects.JaxprInputEffect):
      out_effects.add(e)
      continue
    input_idx = body_input_idx[e.input]
    if input_idx < len(body_jaxpr.constvars):
      const_offset = all_args.body_offset
      out_effects.add(e.replace(const_offset + input_idx))
    else:
      invar_idx = input_idx - len(body_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 should not apply to an input buffer {ref_invar_idx} in"
            f" pipeline body jaxpr: {body_jaxpr}")
      ref_idx = get_ref_idx(flat_refs_idx[ref_invar_idx])
      out_effects.add(e.replace(ref_idx))
  return (), frozenset(out_effects)

# TODO(rdyro): Either generalize or merge with another primitive. This primitive
# perfoms an "eval jaxpr" operation, but is currently tailored to calling the
# pipeline body in the emit_pipeline primtiive - it resolves TransformedRefs and
# binds the user grid indices to lowering.
# This primitive is specialized to resolve TransformedRefs passed as arguments
# and evaluate the body jaxpr with the resolved Refs because it assumes the body
# was traced "generically" with Refs. However, the emit_pipeline is allowed to
# pass in TransformedRefs as arguments to the body.
pipeline_body_p = core.Primitive("pipeline_body")
pipeline_body_p.multiple_results = True

@pipeline_body_p.def_effectful_abstract_eval

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Never assign to in_refs; allocate an out_spec buffer (or scratch) and write results there
2. If in-place semantics are needed, declare that buffer as an output spec with the matching block mapping
3. Check helper functions (e.g. custom update routines) for hidden writes to input refs

Example fix

# before
def body(in_ref, out_ref, i):
  in_ref[i] *= scale  # WRONG: writes input
  out_ref[i] = in_ref[i]
# after
def body(in_ref, out_ref, i):
  out_ref[i] = in_ref[i] * scale

import jax
res = jax.eval_shape(lambda: None)  # ensure no writes: inspect body jaxpr effects in tests
effects = getattr(traced_body, 'effects', None)
assert not any(isinstance(e, jax.experimental.pallas.effects.WriteEffect) for e in (effects or []))

When it happens

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

Common situations: Porting a kernel that used a single scratch ref for both input and output; in-place normalization or caching patterns applied to input buffers.

Related errors

• WriteEffect on input buffer {ref_invar_idx}
• 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/f9b99033a943964e. Report an issue: GitHub.