jax-ml/jax · error · TypeError

reduce_window got the wrong number of window_dimensions for

Error message

reduce_window got the wrong number of window_dimensions for operand: got operand shape {} with window_dimensions {}.

What it means

lax.reduce_window requires len(window_dimensions) == operand.ndim. This shape-rule error fires when the window specification length doesn't match the operand rank.

Source

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

    operand,
    window_dimensions,
    window_strides,
    padding,
    base_dilation,
    window_dilation,
):
  lax._check_shapelike("reduce_window", "window_dimensions", window_dimensions,
                       non_zero_shape=True)
  lax._check_shapelike("reduce_window", "window_strides", window_strides,
                       non_zero_shape=True)
  lax._check_shapelike("reduce_window", "base_dilation", base_dilation)
  lax._check_shapelike("reduce_window", "window_dilation", window_dilation)
  if operand.ndim != len(window_dimensions):
    msg = (
        "reduce_window got the wrong number of window_dimensions for "
        "operand: got operand shape {} with window_dimensions {}."
    )
    raise TypeError(msg.format(operand.shape, window_dimensions))
  if len(window_strides) != len(window_dimensions):
    msg = ("reduce_window got inconsistent window_strides and "
           "window_dimensions: got window_strides {} and window_dimensions {}.")
    raise TypeError(msg.format(window_strides, window_dimensions))
  if len(base_dilation) != len(window_dimensions):
    msg = ("reduce_window got inconsistent base_dilation and "
           "window_dimensions: got base_dilation {} and window_dimensions {}.")
    raise TypeError(msg.format(base_dilation, window_dimensions))
  if len(window_dilation) != len(window_dimensions):
    msg = ("reduce_window got inconsistent window_dilation and "
           "window_dimensions: got window_dilation {} and window_dimensions "
           "{}.")
    raise TypeError(msg.format(window_dilation, window_dimensions))

  return reduce_window_shape_tuple(operand.shape, window_dimensions,
                                   window_strides, padding, base_dilation,
                                   window_dilation)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Extend window_dimensions to include batch and channel dims, usually (1, *spatial, 1) for NHWC
2. Similarly fix window_strides, padding, base_dilation, window_dilation to the same length

Example fix

# before
lax.reduce_window_max(x, (3,3), (1,1), 'VALID')  # x is (N,H,W,C)
# after
lax.reduce_window_max(x, (1,3,3,1), (1,1,1,1), 'VALID')

assert x.ndim == len(window_dimensions), f'rank {x.ndim} vs {len(window_dimensions)} dims'

def window_matches(x, w): return x.ndim == len(w)

When it happens

Trigger: Passing window_dimensions=(3,3) for a 4-D NHWC operand (needs 4 entries like (1,3,3,1)).

Common situations: Converting convolution/pooling code from frameworks expecting only spatial dims; forgetting batch/channel dimensions in NHWC.

Related errors

• reduce_window expected init_values to be scalars but init_va
• reduce_window got inconsistent window_strides and window_dim
• reduce_window got inconsistent base_dilation and window_dime
• reduce_window got inconsistent window_dilation and window_di
• scan got `length` argument of {} which disagrees with leadin

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