jax-ml/jax · error · NotImplementedError

Cannot lower effectful `reduce_window`.

Error message

Cannot lower effectful `reduce_window`.

What it means

When lowering lax.reduce_window to MLIR/XLA, JAX refuses to compile if the reduction body jaxpr has effects (e.g. ordered effects like state, RNG, or IO). XLA's reduce_window requires a pure reducer.

Source

Thrown at jax/_src/lax/windowed_reductions.py:491

def _generic_reduce_window_lower(
    ctx: mlir.LoweringRuleContext,
    *args,
    jaxpr,
    consts,
    window_dimensions,
    window_strides,
    padding,
    base_dilation,
    window_dilation,
):
  operands, init_values = util.split_list(args, [len(args) // 2])
  _, init_value_avals = util.split_list(ctx.avals_in, [len(operands)])

  def reducer_body(reducer: ir.Block) -> Sequence[ir.Value]:
    if jaxpr.effects:
      raise NotImplementedError('Cannot lower effectful `reduce_window`.')
    out_nodes, _ = mlir.jaxpr_subcomp(ctx.module_context, jaxpr, ctx.name_stack,
        mlir.TokenSet(), consts, *reducer.arguments,
        dim_var_values=ctx.dim_var_values, const_lowering=ctx.const_lowering,
        outer_traceback=ctx.traceback)
    flat_out_nodes, _ = mlir.ir_tree_registry.flatten(out_nodes)
    return flat_out_nodes

  return mlir.reduce_window(
      ctx,
      reducer_name="generic_reduce_window_reducer",
      reducer_body=reducer_body,
      operands=operands,
      init_values=init_values,
      init_values_avals=init_value_avals,
      out_avals=ctx.avals_out,
      window_dimensions=window_dimensions,
      window_strides=window_strides,
      base_dilation=base_dilation,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Remove effectful ops (random/state/callback) from the reduction function; keep it pure arithmetic
2. Pre-generate any random values outside the reducer and pass them as consts
3. Use a different decomposition, e.g. scan over windows with explicit effects outside

Example fix

# before
def reducer(x, y):
    return jax.random.bitwise_xor(x, y, key)  # effectful
# after
def reducer(x, y):
    return x + y  # pure

# reducers must be pure; avoid jax.random/state inside
assert 'random' not in jaxpr_effects  # inspect via jax.make_jaxpr(reducer)

try:
    f = jax.jit(pooled)
except NotImplementedError as e:
    if 'effectful' in str(e): rewrite reducer to be pure

When it happens

Trigger: Passing a reduction jaxpr that contains effectful primitives (random bits, stateful ops, host callbacks) to lax.reduce_window.

Common situations: Custom reducers built with jax.random inside; using stateful callbacks in pooling reductions.

Related errors

• Cannot lower jaxpr with effects: {closed_jaxpr.effects}
• Error refining shapes. {dump_module_message(module, "before_
• multi-platform lowering for buffer_callback
• Nesting `compute_on` with different compute types is not all
• Effects not supported in `custom_jvp`: {disallowed_effects}

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