jax-ml/jax · error · NotImplementedError

Cannot lower effectful `reduce_window`.

Error message

Cannot lower effectful `reduce_window`.

What it means

While lowering a reduce_window with a custom reducer jaxpr, the reducer body contains effects. Windowed reduction lowering would need to replicate the effectful body across window positions, which is unsupported, so lowering is refused.

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
• Error refining shapes.
• accessing .backend in multi-lowering setting. This can…
• axis_index_groups must all be the same size for TPU lowering
• cannot cast to

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

Appendix: 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)