{"record":{"id":"d4079ef8e58a867e","repo":"jax-ml/jax","slug":"vjp-not-implemented-for-select-and-gather-maxpool","errorCode":null,"errorMessage":"VJP not implemented for select_and_gather (MaxPool) with window dilation, got window_dilation={}.","messagePattern":"VJP not implemented for select_and_gather \\(MaxPool\\) with window dilation, got window_dilation=(.+?)\\.","errorType":"error_code","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/windowed_reductions.py","lineNumber":1071,"sourceCode":"  del g_operand\n  if type(g_source) is ad_util.Zero:\n    tangent_out = ad_util.p2tz(val_out)\n  else:\n    tangent_out = _select_and_gather_add(\n        g_source, operand, select_prim, window_dimensions,\n        window_strides, padding, base_dilation, window_dilation)\n  return val_out, tangent_out\n\ndef _select_and_gather_add_transpose(\n    t, tangents, operand, *, select_prim, window_dimensions, window_strides,\n    padding, base_dilation, window_dilation):\n  assert select_prim in (lax.le_p, lax.ge_p)\n  assert (ad.is_undefined_primal(tangents) and\n          not ad.is_undefined_primal(operand))\n  if any(d != 1 for d in window_dilation):\n    msg = (\"VJP not implemented for select_and_gather (MaxPool) with window \"\n           \"dilation, got window_dilation={}.\")\n    raise NotImplementedError(msg.format(window_dilation))\n  if type(t) is ad_util.Zero:\n    return [ad_util.Zero(tangents.aval), None]\n  has_base_dilation = any(d != 1 for d in base_dilation)\n  if has_base_dilation:\n    select_identity = (lax._get_max_identity if select_prim is lax.ge_p\n                       else lax._get_min_identity)\n    operand = lax.pad(operand, select_identity(operand.dtype),\n                      tuple((0, 0, d - 1) for d in base_dilation))\n  result = _select_and_scatter_add(t, operand, select_prim, window_dimensions,\n                                   window_strides, padding)\n  if has_base_dilation:\n    result = slicing.slice(result, (0,) * len(result.shape), result.shape,\n                           base_dilation)\n  return [result, None]\n\ndef _select_and_gather_add_batching_rule(\n    batched_args, batch_dims, *, select_prim, window_dimensions, window_strides,\n    padding, base_dilation, window_dilation):","sourceCodeStart":1053,"sourceCodeEnd":1089,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/windowed_reductions.py#L1053-L1089","documentation":"The backward pass (VJP transpose) for max/min pooling (select_and_gather) is not implemented for window dilation != 1. Forward works, but jax.grad fails with NotImplementedError.","triggerScenarios":"Computing gradients of dilated max pooling: lax.reduce_window_max(..., window_dilation=(1,2,2,1)) then jax.grad.","commonSituations":"Differentiating networks with atrous/dilated max-pool layers, e.g. segmentation models.","solutions":["Set window_dilation to 1 in all dims (use strides/downsampling instead of dilated pooling)","Implement dilated pooling manually via slicing/stacking of shifted windows then reduce","Raise an issue upstream / wait for JAX support"],"exampleFix":"# before\npool = lax.reduce_window_max(x, (1,3,3,1), (1,1,1,1), padding, window_dilation=(1,2,2,1))\nloss = pool.sum(); jax.grad(lambda x: loss)(x)\n# after\npool = lax.reduce_window_max(x[:, :, ::2, ::2], (1,3,3,1), (1,1,1,1), padding)\nloss = pool.sum(); jax.grad(lambda x: loss)(x)","handlingStrategy":"fallback","validationCode":"assert all(d == 1 for d in window_dilation), 'VJP unsupported for dilated pooling'","typeGuard":"def vjp_supported(window_dilation): return all(d == 1 for d in window_dilation)","tryCatchPattern":"try:\n    grads = jax.grad(loss)(x)\nexcept NotImplementedError:\n    # fallback: strided slicing instead of dilation\n    grads = jax.grad(lambda x: lax.reduce_window_max(x[:, :, ::2, ::2], (1,3,3,1), (1,1,1,1), 'VALID').sum())(x)","preventionTips":["Avoid dilated max-pool in models needing gradients","Use subsampling/strides to emulate dilation"],"tags":["jax","autodiff","vjp","pooling","dilation","not-implemented"],"backgroundTag":"jax-autodiff-unsupported-differentiation","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}