jax-ml/jax · error · ValueError

Wrong number of pads for spatial dimensions

Error message

Wrong number of pads for spatial dimensions

What it means

The NumPy reference conv/window view requires exactly one (lo, hi) padding pair per spatial dimension (rhs rank minus 2). More or fewer pads raise this ValueError.

Source

Thrown at jax/_src/lax_reference.py:459

      return [
          (pad_size // 2, pad_size - pad_size // 2) for pad_size in pad_sizes
      ]
    else:
      return [
          (pad_size - pad_size // 2, pad_size // 2) for pad_size in pad_sizes
      ]
  else:
    return [(0, 0)] * len(in_shape)

def _conv_view(lhs, rhs_shape, window_strides, pads, pad_value):
  """Compute the view (and its axes) of a convolution or window reduction."""
  if (_min(lhs.ndim, len(rhs_shape)) < 2 or lhs.ndim != len(rhs_shape)
      or lhs.shape[1] != rhs_shape[1]):
    raise ValueError('Dimension mismatch')
  if len(window_strides) != len(rhs_shape) - 2:
    raise ValueError('Wrong number of strides for spatial dimensions')
  if len(pads) != len(rhs_shape) - 2:
    raise ValueError('Wrong number of pads for spatial dimensions')

  lhs = _pad(lhs, [(0, 0)] * 2 + list(pads), pad_value)
  in_shape = lhs.shape[2:]
  filter_shape = rhs_shape[2:]
  dim = len(filter_shape)  # number of 'spatial' dimensions in convolution

  out_strides = np.multiply(window_strides, lhs.strides[2:])
  view_strides = lhs.strides[:1] + tuple(out_strides) + lhs.strides[1:]

  out_shape = np.floor_divide(
      np.subtract(in_shape, filter_shape), window_strides) + 1
  view_shape = lhs.shape[:1] + tuple(out_shape) + rhs_shape[1:]

  view = np.lib.stride_tricks.as_strided(lhs, view_shape, view_strides)

  view_axes = list(range(view.ndim))
  sum_axes = view_axes[-dim-1:]
  rhs_axes = [view.ndim] + sum_axes

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Provide pads as a sequence of (low, high) tuples, one per spatial dimension
2. Convert 'SAME'/'VALID' strings via jax.lax.padtype_to_pads before calling

Example fix

# before
pads = ((0,0),(1,1),(1,1),(0,0))  # full rank
lax_reference._conv_view(..., pads, ...)
# after
pads = ((1,1),(1,1))  # spatial only
lax_reference._conv_view(..., pads, ...)

assert len(pads) == len(rhs_shape) - 2 and all(len(p) == 2 for p in pads)

When it happens

Trigger: Passing padding as full-rank (N+2 pairs) or a flat list of numbers instead of per-spatial-dim pairs.

Common situations: Reusing padding configs from lax.reduce_window (which is full-rank) with the reference conv API (spatial-only).

Related errors

• Wrong number of explicit pads for convolution: expected {},
• Negative padding is larger than the size of the correspondin
• Dimension mismatch
• Wrong number of strides for spatial dimensions
• String padding is not implemented for transposed convolution

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