jax-ml/jax · error · TypeError

select_and_scatter got inconsistent window_strides and windo

Error message

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

What it means

lax.select_and_scatter (used internally for pooling gradients) requires window_strides to have the same length as window_dimensions. This shape rule rejects mismatched lengths.

Source

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

mlir.register_lowering(reduce_window_sum_p, partial(
    _reduce_window_lower, hlo.add, lambda _: 0))
mlir.register_lowering(reduce_window_min_p, partial(
    _reduce_window_lower, mlir.min_hlo, lax._get_min_identity))
mlir.register_lowering(reduce_window_max_p, partial(
    _reduce_window_lower, mlir.max_hlo, lax._get_max_identity))


def _select_and_scatter_shape_rule(
    operand, source, init_value, *, select_jaxpr, select_consts, scatter_jaxpr,
    scatter_consts, window_dimensions, window_strides, padding):
  lax._check_shapelike("select_and_scatter", "window_dimensions",
                       window_dimensions)
  lax._check_shapelike("select_and_scatter", "window_strides", window_strides)
  if len(window_dimensions) != len(window_strides):
    msg = ("select_and_scatter got inconsistent window_strides and "
           "window_dimensions: got window_strides {} and window_dimensions {}.")
    raise TypeError(msg.format(window_strides, window_dimensions))
  return operand.shape

def _select_and_scatter_sharding_rule(
    operand, source, init_value, *, select_jaxpr, select_consts, scatter_jaxpr,
    scatter_consts, window_dimensions, window_strides, padding):
  return operand.sharding

select_and_scatter_p = lax.standard_primitive(
    _select_and_scatter_shape_rule, lax.input_dtype, 'select_and_scatter',
    sharding_rule=_select_and_scatter_sharding_rule,
    vma_rule=partial(core.standard_vma_rule, 'select_and_scatter'))

def _select_and_scatter_lower(
    ctx: mlir.LoweringRuleContext, operand, source, init_value, *,
    select_jaxpr: core.Jaxpr, select_consts,
    scatter_jaxpr: core.Jaxpr, scatter_consts, window_dimensions,
    window_strides, padding):
  operand_aval, source_aval, init_value_aval = ctx.avals_in

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Match lengths: give one stride per dimension including batch/channel
2. Prefer higher-level jax.lax.reduce_window or NN library pooling instead of manual select_and_scatter

Example fix

# before
lax.select_and_scatter(operand, source, init, ..., (2,2), (1,2,2,1), padding)
# after
lax.select_and_scatter(operand, source, init, ..., (1,2,2,1), (1,2,2,1), padding)

assert len(window_strides) == len(window_dimensions)

When it happens

Trigger: Directly calling lax.select_and_scatter with len(window_strides) != len(window_dimensions).

Common situations: Writing custom pooling backward passes or porting an XLA SelectAndScatter computation.

Related errors

• select_and_gather_add tangents and operand shapes must match
• scan got `length` argument of {} which disagrees with leadin
• conv_general_dilated batch_group_count must divide lhs batch
• conv_general_dilated rhs output feature dimension size must
• conv_general_dilated window and window_strides must have the

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