jax-ml/jax · error · TypeError

reduce_window got inconsistent window_dilation and window_di

Error message

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

What it means

lax.reduce_window's window_dilation must match window_dimensions in length. Raised when len(window_dilation) != len(window_dimensions).

Source

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

  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:
    operand_shape = lax._dilate_shape(operand_shape, base_dilation)
  if window_dilation is not None:
    window_dimensions = lax._dilate_shape(window_dimensions, window_dilation)
  operand_padded = tuple(d + pl + ph for d, (pl, ph) in zip(operand_shape, padding))
  return tuple(map(core.stride_dim, operand_padded, window_dimensions, window_strides))


reduce_window_max_p = lax.standard_primitive(
    _common_reduce_window_shape_rule, lax.input_dtype, 'reduce_window_max',

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Extend window_dilation to full rank, e.g. (1, dh, dw, 1) for NHWC operands

Example fix

# before
lax.reduce_window_max(x, (1,3,3,1), (1,1,1,1), padding, window_dilation=(2,2))
# after
lax.reduce_window_max(x, (1,3,3,1), (1,1,1,1), padding, window_dilation=(1,2,2,1))

Defensive patterns

Strategy: validation

Validate before calling

assert len(window_dilation) == len(window_dimensions)

Prevention

Always pass full-rank dilation tuples in NHWC order

When it happens

Trigger: Passing window_dilation=(1,1) or defaulting while window_dimensions is length 4 (NHWC).

Common situations: Implementing atrous/dilated max-pooling and giving only spatial dilation values.

Related errors

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