{"record":{"id":"ebfc03bd4c839336","repo":"jax-ml/jax","slug":"select-and-scatter-got-inconsistent-window-strides","errorCode":null,"errorMessage":"select_and_scatter got inconsistent window_strides and window_dimensions: got window_strides {} and window_dimensions {}.","messagePattern":"select_and_scatter got inconsistent window_strides and window_dimensions: got window_strides (.+?) and window_dimensions (.+?)\\.","errorType":"validation","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/windowed_reductions.py","lineNumber":713,"sourceCode":"\nmlir.register_lowering(reduce_window_sum_p, partial(\n    _reduce_window_lower, hlo.add, lambda _: 0))\nmlir.register_lowering(reduce_window_min_p, partial(\n    _reduce_window_lower, mlir.min_hlo, lax._get_min_identity))\nmlir.register_lowering(reduce_window_max_p, partial(\n    _reduce_window_lower, mlir.max_hlo, lax._get_max_identity))\n\n\ndef _select_and_scatter_shape_rule(\n    operand, source, init_value, *, select_jaxpr, select_consts, scatter_jaxpr,\n    scatter_consts, window_dimensions, window_strides, padding):\n  lax._check_shapelike(\"select_and_scatter\", \"window_dimensions\",\n                       window_dimensions)\n  lax._check_shapelike(\"select_and_scatter\", \"window_strides\", window_strides)\n  if len(window_dimensions) != len(window_strides):\n    msg = (\"select_and_scatter got inconsistent window_strides and \"\n           \"window_dimensions: got window_strides {} and window_dimensions {}.\")\n    raise TypeError(msg.format(window_strides, window_dimensions))\n  return operand.shape\n\ndef _select_and_scatter_sharding_rule(\n    operand, source, init_value, *, select_jaxpr, select_consts, scatter_jaxpr,\n    scatter_consts, window_dimensions, window_strides, padding):\n  return operand.sharding\n\nselect_and_scatter_p = lax.standard_primitive(\n    _select_and_scatter_shape_rule, lax.input_dtype, 'select_and_scatter',\n    sharding_rule=_select_and_scatter_sharding_rule,\n    vma_rule=partial(core.standard_vma_rule, 'select_and_scatter'))\n\ndef _select_and_scatter_lower(\n    ctx: mlir.LoweringRuleContext, operand, source, init_value, *,\n    select_jaxpr: core.Jaxpr, select_consts,\n    scatter_jaxpr: core.Jaxpr, scatter_consts, window_dimensions,\n    window_strides, padding):\n  operand_aval, source_aval, init_value_aval = ctx.avals_in","sourceCodeStart":695,"sourceCodeEnd":731,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/windowed_reductions.py#L695-L731","documentation":"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.","triggerScenarios":"Directly calling lax.select_and_scatter with len(window_strides) != len(window_dimensions).","commonSituations":"Writing custom pooling backward passes or porting an XLA SelectAndScatter computation.","solutions":["Match lengths: give one stride per dimension including batch/channel","Prefer higher-level jax.lax.reduce_window or NN library pooling instead of manual select_and_scatter"],"exampleFix":"# before\nlax.select_and_scatter(operand, source, init, ..., (2,2), (1,2,2,1), padding)\n# after\nlax.select_and_scatter(operand, source, init, ..., (1,2,2,1), (1,2,2,1), padding)","handlingStrategy":"validation","validationCode":"assert len(window_strides) == len(window_dimensions)","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Prefer jax.lax.reduce_window over raw select_and_scatter"],"tags":["jax","shape-validation","select-and-scatter","pooling"],"backgroundTag":"jax-pooling-gradient-invalid-arguments","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}