jax-ml/jax · error · TypeError

select_and_gather_add tangents and operand shapes must match

Error message

select_and_gather_add tangents and operand shapes must match, got {} and {}.

What it means

The VJP transpose rule for select_and_gather (the max/min pooling gradient) requires the cotangent (tangents) array to have exactly the same shape as the operand. This internal rule fires during jax.grad of pooling when shapes drift.

Source

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

    partial(_select_and_scatter_add_impl, expand_padding=False),
    multiple_results=False))
# TODO(b/161704903): workaround for XLA/CPU crash.
mlir.register_lowering(select_and_scatter_add_p, mlir.lower_fun(
    partial(_select_and_scatter_add_impl, expand_padding=True),
    multiple_results=False), platform='cpu')
# TODO(b/182390722): workaround for XLA/GPU crash.
mlir.register_lowering(select_and_scatter_add_p, mlir.lower_fun(
    partial(_select_and_scatter_add_impl, expand_padding=True),
    multiple_results=False), platform='gpu')


def _select_and_gather_add_shape_rule(
    tangents, operand, *, select_prim, window_dimensions, window_strides,
    padding, base_dilation, window_dilation):
  if tangents.shape != operand.shape:
    msg = ("select_and_gather_add tangents and operand shapes must match, "
           "got {} and {}.")
    raise TypeError(msg.format(tangents.shape, operand.shape))
  return _common_reduce_window_shape_rule(
      operand, window_dimensions, window_strides, padding, base_dilation,
      window_dilation)

def _select_and_gather_add_sharding_rule(
    tangents, operand, *, select_prim, window_dimensions, window_strides,
    padding, base_dilation, window_dilation):
  if tangents.sharding != operand.sharding:
    raise core.ShardingTypeError(
        "select_and_gather_add tangents and operand shardings must match, "
        f"got {tangents.sharding} and {operand.sharding}.")
  return reduce_window_sharding_rule(
      operand, window_dimensions, window_strides, padding, base_dilation,
      window_dilation)

def _select_and_gather_add_lowering(
    ctx: mlir.LoweringRuleContext,
    tangents, operand, *, select_prim,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Ensure the cotangent produced matches the operand shape; avoid reshapes between pooling and downstream ops
2. Use standard pooling APIs (e.g. jax.lax.reduce_window_max) whose transpose rules keep shapes aligned
3. Check padding config so forward output shape matches what the loss produces

assert tangents.shape == operand.shape

When it happens

Trigger: Applying jax.grad through max_pool where the incoming cotangent shape differs from the pooled input's shape (e.g. custom pooling wrapper reshapes between pool and loss).

Common situations: Custom pooling implementations that reshape or slice the operand inside the differentiable path; mismatched padding making output shapes inconsistent.

Related errors

• VJP not implemented for select_and_gather (MaxPool) with win
• the gradient for {_vjp_arg_name(jaxpr, in_tree, idx)}, which
• {_vjp_arg_name(jaxpr, in_tree, idx)} is Ref-typed, so its gr
• structure of the differentiated function {jaxpr.debug_info.f
• Pure callbacks do not support transpose. Please use `jax.cus

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