{"record":{"id":"13ab9b9ed026ccb2","repo":"jax-ml/jax","slug":"reduce-window-batching-is-not-implemented-for-init","errorCode":null,"errorMessage":"reduce_window batching is not implemented for initial values","messagePattern":"reduce_window batching is not implemented for initial values","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/windowed_reductions.py","lineNumber":403,"sourceCode":"\ndef _generic_reduce_window_batch_rule(\n    batched_args,\n    batch_dims,\n    *,\n    jaxpr,\n    consts,\n    window_dimensions,\n    window_strides,\n    padding,\n    base_dilation,\n    window_dilation,\n):\n  num_operands = len(batched_args) // 2\n  operands, init_values = util.split_list(batched_args, [num_operands])\n  operand_bdims, init_value_bdims = util.split_list(batch_dims, [num_operands])\n\n  if any(init_bdim is not None for init_bdim in init_value_bdims):\n    raise NotImplementedError(\"reduce_window batching is not implemented for \"\n                              \"initial values\")\n\n  size = next(x.shape[ax] for x, ax in zip(operands, operand_bdims)\n              if ax is not None)\n  operands = [batching.bdim_at_front(arg, bdim, size)\n              for arg, bdim in zip(operands, operand_bdims)]\n  window_dimensions = (1,) + window_dimensions\n  window_strides = (1,) + window_strides\n  padding = ((0, 0),) + padding\n  base_dilation = (1,) + base_dilation\n  window_dilation = (1,) + window_dilation\n  outs = reduce_window_p.bind(\n      *(operands + init_values), jaxpr=jaxpr, consts=consts,\n      window_dimensions=window_dimensions, window_strides=window_strides,\n      padding=padding, base_dilation=base_dilation,\n      window_dilation=window_dilation)\n  return outs, (0,) * num_operands\n","sourceCodeStart":385,"sourceCodeEnd":421,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/windowed_reductions.py#L385-L421","documentation":"JAX's vmap batching rule for the generic lax.reduce_window does not support batching over the init_value arguments. When vmap maps an axis that flows into an init_value, this NotImplementedError is raised.","triggerScenarios":"jax.vmap over a function whose reduce_window call receives an init_value that depends on the batched argument (e.g. per-sample identity values of shape (batch,)).","commonSituations":"Batching pooling layers or custom reductions where identity values vary per batch element.","solutions":["Make init_value a constant scalar independent of the batched input (e.g. -inf for max)","If per-sample init values are essential, implement the reduction manually with jax.lax.scan or jnp ops","Use jax.vmap with in_axes that exclude the init_value argument"],"exampleFix":"// before\njax.vmap(lambda x, iv: lax.reduce_window(x, iv, ...))(xs, inits)\n// after\njax.vmap(lambda x: lax.reduce_window(x, -np.inf, ...))(xs)","handlingStrategy":"validation","validationCode":"# keep init_value independent of batched inputs\ninit = jnp.asarray(-np.inf)  # constant, never shape (batch,)","typeGuard":null,"tryCatchPattern":"try:\n    jax.vmap(fn)(xs)\nexcept NotImplementedError as e:\n    if 'batching is not implemented' in str(e):\n        fn = lambda x: lax.reduce_window(x, -np.inf, ...)  # fixed init","preventionTips":["Never derive init_value from batched data inside vmap","Test vmapped pooling functions in CI"],"tags":["jax","vmap","batching","not-implemented","reduce-window"],"backgroundTag":"jax-vmap-unsupported-operation","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}