jax-ml/jax · error · TypeError

select_and_gather_add tangents and operand shapes must…

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.

Solutions

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

Defensive patterns

Strategy: validation

Validate before calling

assert tangents.shape == operand.shape

Prevention

Avoid reshaping between pooling and loss in differentiable paths
Use standard pooling APIs

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

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

Appendix: 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)