jax-ml/jax · error · NotImplementedError

VJP not implemented for select_and_gather (MaxPool) with win

Error message

VJP not implemented for select_and_gather (MaxPool) with window dilation, got window_dilation={}.

What it means

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.

Source

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

  del g_operand
  if type(g_source) is ad_util.Zero:
    tangent_out = ad_util.p2tz(val_out)
  else:
    tangent_out = _select_and_gather_add(
        g_source, operand, select_prim, window_dimensions,
        window_strides, padding, base_dilation, window_dilation)
  return val_out, tangent_out

def _select_and_gather_add_transpose(
    t, tangents, operand, *, select_prim, window_dimensions, window_strides,
    padding, base_dilation, window_dilation):
  assert select_prim in (lax.le_p, lax.ge_p)
  assert (ad.is_undefined_primal(tangents) and
          not ad.is_undefined_primal(operand))
  if any(d != 1 for d in window_dilation):
    msg = ("VJP not implemented for select_and_gather (MaxPool) with window "
           "dilation, got window_dilation={}.")
    raise NotImplementedError(msg.format(window_dilation))
  if type(t) is ad_util.Zero:
    return [ad_util.Zero(tangents.aval), None]
  has_base_dilation = any(d != 1 for d in base_dilation)
  if has_base_dilation:
    select_identity = (lax._get_max_identity if select_prim is lax.ge_p
                       else lax._get_min_identity)
    operand = lax.pad(operand, select_identity(operand.dtype),
                      tuple((0, 0, d - 1) for d in base_dilation))
  result = _select_and_scatter_add(t, operand, select_prim, window_dimensions,
                                   window_strides, padding)
  if has_base_dilation:
    result = slicing.slice(result, (0,) * len(result.shape), result.shape,
                           base_dilation)
  return [result, None]

def _select_and_gather_add_batching_rule(
    batched_args, batch_dims, *, select_prim, window_dimensions, window_strides,
    padding, base_dilation, window_dilation):

View on GitHub (pinned to 1e1c6a8fc0)

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

Example fix

# before
pool = lax.reduce_window_max(x, (1,3,3,1), (1,1,1,1), padding, window_dilation=(1,2,2,1))
loss = pool.sum(); jax.grad(lambda x: loss)(x)
# after
pool = lax.reduce_window_max(x[:, :, ::2, ::2], (1,3,3,1), (1,1,1,1), padding)
loss = pool.sum(); jax.grad(lambda x: loss)(x)

Defensive patterns

Strategy: fallback

Validate before calling

assert all(d == 1 for d in window_dilation), 'VJP unsupported for dilated pooling'

Type guard

def vjp_supported(window_dilation): return all(d == 1 for d in window_dilation)

Try / catch

try:
    grads = jax.grad(loss)(x)
except NotImplementedError:
    # fallback: strided slicing instead of dilation
    grads = jax.grad(lambda x: lax.reduce_window_max(x[:, :, ::2, ::2], (1,3,3,1), (1,1,1,1), 'VALID').sum())(x)

Prevention

Avoid dilated max-pool in models needing gradients
Use subsampling/strides to emulate dilation

When it happens

Trigger: Computing gradients of dilated max pooling: lax.reduce_window_max(..., window_dilation=(1,2,2,1)) then jax.grad.

Common situations: Differentiating networks with atrous/dilated max-pool layers, e.g. segmentation models.

Related errors

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