jax-ml/jax · error · TypeError

Wrong number of explicit pads for convolution: expected {},

Error message

Wrong number of explicit pads for convolution: expected {}, got {}.

What it means

conv_shape_tuple computes the output shape of a convolution from canonical (N,H,W,C) input shapes. When padding is given as an explicit list of (lo,hi) pairs, there must be exactly one pair per spatial dimension (lhs rank minus 2). Any other length raises this TypeError.

Source

Thrown at jax/_src/lax/convolution.py:904

  new_shape = list(np.delete(x.shape, src))
  new_shape[dst] *= x.shape[src]
  return lax.reshape(x, new_shape, perm)

def _reshape_axis_out_of(src, size1, x):
  shape = list(x.shape)
  size2, ragged = divmod(shape[src], size1)
  assert not ragged
  shape[src:src+1] = [size1, size2]
  return lax.reshape(x, shape)


def conv_shape_tuple(lhs_shape, rhs_shape, strides, pads, batch_group_count=1):
  """Compute the shape tuple of a conv given input shapes in canonical order."""
  if isinstance(pads, str):
    pads = lax.padtype_to_pads(lhs_shape[2:], rhs_shape[2:], strides, pads)
  if len(pads) != len(lhs_shape) - 2:
    msg = "Wrong number of explicit pads for convolution: expected {}, got {}."
    raise TypeError(msg.format(len(lhs_shape) - 2, len(pads)))

  lhs_padded = np.add(lhs_shape[2:], np.sum(np.array(pads).reshape(-1, 2),
                                              axis=1))
  if np.any(lhs_padded < 0):
    raise ValueError("Negative padding is larger than the size of the corresponding dimension: "
                     f"got padding={pads} for lhs_shape[2:]={lhs_shape[2:]}")
  out_space = tuple(map(core.stride_dim, lhs_padded, rhs_shape[2:], strides))
  if batch_group_count > 1:
    assert lhs_shape[0] % batch_group_count == 0
    out_shape_0 = lhs_shape[0] // batch_group_count
  else:
    out_shape_0 = lhs_shape[0]
  out_shape = (out_shape_0, rhs_shape[0])
  return tuple(out_shape + tuple(out_space))


def conv_general_shape_tuple(lhs_shape, rhs_shape, window_strides, padding,
                             dimension_numbers):

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Provide pads as a sequence of (low, high) pairs, one per spatial dim: ((0,0),(1,1)) for 2D
2. Or pass padding as a string 'SAME'/'VALID' and let JAX compute pads
3. Double-check the spatial rank: len(pads) must equal len(lhs_shape)-2

Example fix

# before
out = lax.conv_general_dilated(x, w, (1,1), [(0,0,1,1)])
# after
out = lax.conv_general_dilated(x, w, (1,1), ((0,0),(1,1)))

assert len(pads) == len(lhs.shape) - 2 and all(len(p) == 2 for p in pads), f'pads {pads} for rank {lhs.ndim}'

def valid_pads(pads, spatial_ndim) -> bool:
    return (isinstance(pads, str) and pads in ('SAME','VALID')) or (
        isinstance(pads, (list, tuple)) and len(pads) == spatial_ndim and
        all(len(p) == 2 for p in pads))

When it happens

Trigger: Passing pads as e.g. [(0,0)] for a 2-spatial-dim conv, or passing a flat list of 4 ints instead of 2 pairs, to functions like lax.conv_general_dilated / conv_general_shape_tuple with explicit pads.

Common situations: Hand-constructing padding lists after reading XLA conv specs; converting 'SAME'/'VALID' logic into explicit pads incorrectly; mixing up rank between 1D/2D convs.

Related errors

• Negative padding is larger than the size of the correspondin
• Wrong number of pads for spatial dimensions
• String padding is not implemented for transposed convolution
• padding argument to conv_general_dilated should be a string
• conv_general_dilated batch_group_count must divide lhs batch

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