jax-ml/jax · error · ValueError

Parameter {key} is not a Jaxpr or sequence of Jaxprs: {value

Error message

Parameter {key} is not a Jaxpr or sequence of Jaxprs: {value}

What it means

This error is thrown by the Pallas HLO interpreter when a higher-order primitive (scan, while_loop, cond) is being re-bound with parameters that were expected to contain a Jaxpr or a sequence of Jaxprs, but the parameter value is neither. The interpreter's rule factory (make_hop_rule) inspects each primitive's params and must lower embedded jaxprs to interpret them; unknown parameter shapes abort.

Source

Thrown at jax/_src/pallas/hlo_interpreter.py:275

    new_jaxpr = jaxpr.with_consts(new_consts)
    return new_jaxpr, extra_args

  def rule(interpreter, *args, **params):
    new_params = {}
    for key in keys:
      value = params[key]
      if isinstance(value, jax_core.Jaxpr):
        new_jaxpr, extra_args = _resolve_jaxpr(interpreter, value)
        new_params[key] = new_jaxpr
        args = extra_args + args
      elif isinstance(value, tuple) or isinstance(value, list):
        mapped_jaxprs, mapped_args = zip(*map(
          lambda x, i: _resolve_jaxpr(interpreter, x, mapped_idx=i), value, range(len(value))))
        all_new_args = tuple(new_arg for _args in mapped_args for new_arg in _args)
        new_params[key] = tuple(mapped_jaxprs)
        args = all_new_args + args
      else:
        raise ValueError(f"Parameter {key} is not a Jaxpr or sequence of Jaxprs: {value}")
    params.update(new_params)
    return primitive.bind(*args, **params)
  return rule

_eval_jaxpr_hop_rules[loops.scan_p] = make_hop_rule(loops.scan_p, 'jaxpr')
_eval_jaxpr_hop_rules[loops.while_p] = make_hop_rule(
    loops.while_p, 'body_jaxpr', 'cond_jaxpr')
_eval_jaxpr_hop_rules[conditionals.cond_p] = make_hop_rule(conditionals.cond_p, 'branches')
def _run_scoped_physicalize_rule(
    interpreter, *consts, jaxpr: jax_core.Jaxpr, collective_axes, **params):
  if collective_axes:
    raise NotImplementedError(
        "run_scoped interpret rule does not support collective axes"
    )
  physical_jaxpr, physical_consts = interpreter(jaxpr, consts)
  return primitives.run_scoped_p.bind(
      *physical_consts, jaxpr=physical_jaxpr, collective_axes=collective_axes,
      **params

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Ensure the parameter passed for the control-flow primitive is a raw jax_core.Jaxpr, not a ClosedJaxpr or tuple — unwrap with closed_jaxpr.jaxpr and thread consts as args
2. Check that any custom primitive params added by your transformation don't shadow the jaxpr-typed keys ('jaxpr', 'body_jaxpr', 'cond_jaxpr', 'branches')
3. Update JAX to a matching version where the interpreter rule supports your primitive's parameter layout
4. Avoid the HLO interpreter path (compile normally or use the standard interpret mode) if your kernel uses exotic control-flow metadata

Example fix

// before
params['body_jaxpr'] = (closed_jaxpr, ())  # tuple, not a Jaxpr
// after
params['body_jaxpr'] = closed_jaxpr.jaxpr
args = args + tuple(closed_jaxpr.consts)

import jax._src.core as jax_core
def is_jaxpr(v):
    return isinstance(v, jax_core.Jaxpr)

def is_jaxpr(v) -> TypeGuard[jax_core.Jaxpr]: return isinstance(v, jax_core.Jaxpr)

When it happens

Trigger: Running the Pallas HLO interpreter (e.g., pallas_call with an interpret/translation stage that hits _eval_jaxpr_hop_rules) over a kernel containing scan_p, while_p, or cond_p whose jaxpr-typed parameter ('jaxpr', 'body_jaxpr', 'cond_jaxpr', 'branches') holds a non-Jaxpr value (e.g., a closed jaxpr tuple, None, or a custom params dict injected by a transformation).

Common situations: Custom JAX transformations or new primitives that attach extra params to control-flow primitives; version mismatches where a param's expected type changed; passing a jaxpr closure via pjit-style (closed_jaxpr, consts) tuples instead of a raw Jaxpr.

Related errors

• run_scoped interpret rule does not support collective axes
• Padding with aliasing not supported.
• Advanced indexers are not supported on TPU
• Revisited block {output_ranges[i]} of output {i} in iteratio
• Scan with {num_extensive_outputs} extensive output(s) is not

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