jax-ml/jax · error · NotImplementedError

reduce_window batching is not implemented for initial values

Error message

reduce_window batching is not implemented for initial values

What it means

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.

Source

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

def _generic_reduce_window_batch_rule(
    batched_args,
    batch_dims,
    *,
    jaxpr,
    consts,
    window_dimensions,
    window_strides,
    padding,
    base_dilation,
    window_dilation,
):
  num_operands = len(batched_args) // 2
  operands, init_values = util.split_list(batched_args, [num_operands])
  operand_bdims, init_value_bdims = util.split_list(batch_dims, [num_operands])

  if any(init_bdim is not None for init_bdim in init_value_bdims):
    raise NotImplementedError("reduce_window batching is not implemented for "
                              "initial values")

  size = next(x.shape[ax] for x, ax in zip(operands, operand_bdims)
              if ax is not None)
  operands = [batching.bdim_at_front(arg, bdim, size)
              for arg, bdim in zip(operands, operand_bdims)]
  window_dimensions = (1,) + window_dimensions
  window_strides = (1,) + window_strides
  padding = ((0, 0),) + padding
  base_dilation = (1,) + base_dilation
  window_dilation = (1,) + window_dilation
  outs = reduce_window_p.bind(
      *(operands + init_values), jaxpr=jaxpr, consts=consts,
      window_dimensions=window_dimensions, window_strides=window_strides,
      padding=padding, base_dilation=base_dilation,
      window_dilation=window_dilation)
  return outs, (0,) * num_operands

View on GitHub (pinned to 1e1c6a8fc0)

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

Example fix

// before
jax.vmap(lambda x, iv: lax.reduce_window(x, iv, ...))(xs, inits)
// after
jax.vmap(lambda x: lax.reduce_window(x, -np.inf, ...))(xs)

Defensive patterns

Strategy: validation

Validate before calling

# keep init_value independent of batched inputs
init = jnp.asarray(-np.inf)  # constant, never shape (batch,)

Try / catch

try:
    jax.vmap(fn)(xs)
except NotImplementedError as e:
    if 'batching is not implemented' in str(e):
        fn = lambda x: lax.reduce_window(x, -np.inf, ...)  # fixed init

Prevention

Never derive init_value from batched data inside vmap
Test vmapped pooling functions in CI

When it happens

Trigger: 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,)).

Common situations: Batching pooling layers or custom reductions where identity values vary per batch element.

Related errors

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