{"record":{"id":"2d744eb0435dfc7c","repo":"jax-ml/jax","slug":"cannot-lower-effectful-reduce-window","errorCode":null,"errorMessage":"Cannot lower effectful `reduce_window`.","messagePattern":"Cannot lower effectful `reduce_window`\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/windowed_reductions.py","lineNumber":491,"sourceCode":"\n\ndef _generic_reduce_window_lower(\n    ctx: mlir.LoweringRuleContext,\n    *args,\n    jaxpr,\n    consts,\n    window_dimensions,\n    window_strides,\n    padding,\n    base_dilation,\n    window_dilation,\n):\n  operands, init_values = util.split_list(args, [len(args) // 2])\n  _, init_value_avals = util.split_list(ctx.avals_in, [len(operands)])\n\n  def reducer_body(reducer: ir.Block) -> Sequence[ir.Value]:\n    if jaxpr.effects:\n      raise NotImplementedError('Cannot lower effectful `reduce_window`.')\n    out_nodes, _ = mlir.jaxpr_subcomp(ctx.module_context, jaxpr, ctx.name_stack,\n        mlir.TokenSet(), consts, *reducer.arguments,\n        dim_var_values=ctx.dim_var_values, const_lowering=ctx.const_lowering,\n        outer_traceback=ctx.traceback)\n    flat_out_nodes, _ = mlir.ir_tree_registry.flatten(out_nodes)\n    return flat_out_nodes\n\n  return mlir.reduce_window(\n      ctx,\n      reducer_name=\"generic_reduce_window_reducer\",\n      reducer_body=reducer_body,\n      operands=operands,\n      init_values=init_values,\n      init_values_avals=init_value_avals,\n      out_avals=ctx.avals_out,\n      window_dimensions=window_dimensions,\n      window_strides=window_strides,\n      base_dilation=base_dilation,","sourceCodeStart":473,"sourceCodeEnd":509,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/windowed_reductions.py#L473-L509","documentation":"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.","triggerScenarios":"Passing a reduction jaxpr that contains effectful primitives (random bits, stateful ops, host callbacks) to lax.reduce_window.","commonSituations":"Custom reducers built with jax.random inside; using stateful callbacks in pooling reductions.","solutions":["Remove effectful ops (random/state/callback) from the reduction function; keep it pure arithmetic","Pre-generate any random values outside the reducer and pass them as consts","Use a different decomposition, e.g. scan over windows with explicit effects outside"],"exampleFix":"# before\ndef reducer(x, y):\n    return jax.random.bitwise_xor(x, y, key)  # effectful\n# after\ndef reducer(x, y):\n    return x + y  # pure","handlingStrategy":"validation","validationCode":"# reducers must be pure; avoid jax.random/state inside\nassert 'random' not in jaxpr_effects  # inspect via jax.make_jaxpr(reducer)","typeGuard":null,"tryCatchPattern":"try:\n    f = jax.jit(pooled)\nexcept NotImplementedError as e:\n    if 'effectful' in str(e): rewrite reducer to be pure","preventionTips":["Keep reduction bodies to pure arithmetic","Precompute random constants outside the reducer"],"tags":["jax","effects","lowering","mlir","reduce-window"],"backgroundTag":"jax-effectful-primitive-lowering","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}