jax-ml/jax · error · TypeError

reduce_window got inconsistent window_strides and window_dim

Error message

reduce_window got inconsistent window_strides and window_dimensions: got window_strides {} and window_dimensions {}.

What it means

All windowing parameters of lax.reduce_window must have the same length as window_dimensions. This fires when len(window_strides) != len(window_dimensions).

Source

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

    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)

def reduce_window_shape_tuple(operand_shape, window_dimensions, window_strides,
                              padding, base_dilation=None,
                              window_dilation=None):
  if base_dilation is not None:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Make window_strides the same length as window_dimensions, padding to 1 for non-spatial axes

Example fix

# before
lax.reduce_window(x, init, jaxpr, (1,3,3,1), (2,2), padding)
# after
lax.reduce_window(x, init, jaxpr, (1,3,3,1), (1,2,2,1), padding)

Defensive patterns

Strategy: validation

Validate before calling

assert len(window_strides) == len(window_dimensions)

Prevention

Build all window params from one spatial tuple, e.g. full=lambda s: (1, *s, 1)

When it happens

Trigger: Passing window_strides=(1,1) with window_dimensions=(1,3,3,1) (mismatched lengths).

Common situations: Building pooling layers where strides list was updated but window dims weren't (or vice versa).

Related errors

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