jax-ml/jax · error · TypeError

reduce_window got inconsistent base_dilation and window_dime

Error message

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

What it means

lax.reduce_window's base_dilation must have the same length as window_dimensions. Raised by the common shape rule when len(base_dilation) mismatches.

Source

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

                       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:
    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))

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Provide base_dilation of the same rank as the operand, e.g. (1,dh,dw,1) in NHWC

Example fix

# before
lax.reduce_window_max(x, dims, strides, padding, base_dilation=(2,2))
# after
lax.reduce_window_max(x, dims, strides, padding, base_dilation=(1,2,2,1))

Defensive patterns

Strategy: validation

Validate before calling

assert len(base_dilation) == len(window_dimensions)

Prevention

Always pass full-rank dilation tuples in NHWC order

When it happens

Trigger: Supplying base_dilation=(1,1) with a 4-element window_dimensions; or relying on the default while passing mis-sized lists.

Common situations: Using dilated pooling (e.g. mixed-conv downsampling) with incomplete dilation specs.

Related errors

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