jax-ml/jax · error · ValueError

`JaxprInputEffect` {eff} does not have a corresponding jaxpr

Error message

`JaxprInputEffect` {eff} does not have a corresponding jaxpr input.
 Equation: {eqn}

 Effects: {eqn.effects}

 Jaxpr: {core.Jaxpr(constvars, invars, outvars, eqns, set(), dbg)}

What it means

When a JaxprEqn carries a JaxprInputEffect (an effect tied to a specific jaxpr input, e.g. a mutable-array get/put effect referencing input #2), the referenced input must actually be one of the jaxpr's input vars. If the effect's input index points at a var that isn't in input_vars, jaxpr construction is inconsistent and JAX raises with the full eqn, effects, and jaxpr printed.

Source

Thrown at jax/_src/interpreters/partial_eval.py:1463

def make_jaxpr_effects(constvars, invars, outvars, eqns) -> effects.Effects:
  jaxpr_effects = set()
  input_vars = {*constvars, *invars}
  mut_arrays = set()
  for eqn in eqns:
    if eqn.primitive.ref_allocating:
      outvar, = eqn.outvars
      mut_arrays.add(outvar)
    for eff in eqn.effects:
      if isinstance(eff, effects.JaxprInputEffect):
        if eff.input in mut_arrays:
          continue
        if eff.input not in input_vars:
          # TODO(mattjj): ask for forgiveness
          dbg = type('Fake', (), {'resolve_result_paths': lambda self_: self_,
                                  'assert_arg_names': lambda _, __: None,
                                  'assert_result_paths': lambda _, __: None,
                                  })()
          raise ValueError(
                f"`JaxprInputEffect` {eff} does not have "
                f"a corresponding jaxpr input."
                f"\n Equation: {eqn}\n"
                f"\n Effects: {eqn.effects}\n"
                "\n Jaxpr: "
                f"{core.Jaxpr(constvars, invars, outvars, eqns, set(), dbg)}")
      jaxpr_effects.add(eff)
  return jaxpr_effects


class JaxprStackFrame:
  __slots__ = (
      'gensym', 'constid_to_tracer', 'constvar_to_val', 'tracing_eqns',
      'invars', 'effects', 'debug_info', 'is_high', 'auto_dce')

  gensym: Callable[[AbstractValue], Var]
  constid_to_tracer: WeakValueDictionary[ConstId, DynamicJaxprTracer]
  constvar_to_val: dict[Var, Any]

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Update/downgrade JAX to a consistent version — many instances are fixed internal bugs
2. Reduce the reproducer (usually remat + mutable state + DCE interplay) and report it upstream with the printed jaxpr
3. Avoid combining experimental Ref/mutable-array state with remat or heavy DCE-triggering code until fixed
4. If you build eqns manually, ensure JaxprInputEffect.input indices stay within the surviving invars

try:
    jaxpr = jax.make_jaxpr(f)(*args)
except ValueError as e:
    if 'JaxprInputEffect' in str(e):
        # likely internal bug: capture full message (contains jaxpr dump) and report

When it happens

Trigger: Usually an internal invariant failure or a bug in custom primitives/effects that construct eqns manually; can surface via tracers_to_jaxpr, _dce_jaxpr (dead-code elimination removing an input an effect still references), or partial-eval custom caching (_partial_eval_jaxpr_custom_cached).

Common situations: JAX version mismatches after upgrades; custom primitives with declared input effects where DCE or partial evaluation removes the referenced input; bugs in experimental mutable-array code paths.

Related errors

• Effects not supported in `custom_jvp`: {disallowed_effects}
• Effects not supported in `custom_jvp`: {disallowed}
• Cannot lower jaxpr with effects: {closed_jaxpr.effects}
• function {dbg.func_src_info} traced for {dbg.traced_for} ret
• The following ordered effects are not supported for more tha

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