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.

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)))
Defensive patterns

Strategy: validation

Validate before calling

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

Type guard

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

Prevention

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.

Understand the failure class

Background: "Invalid ... format", "must be in format X", "does not look like a ..." — invalid argument format errors across CLI tools and libraries — this error's family across 17 libraries.

Related errors

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

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