jax-ml/jax · error · ValueError

Invalid padding mode: {padding}

Error message

Invalid padding mode: {padding}

What it means

conv_transpose derives the effective padding for its underlying dilated convolution from the requested padding. Only 'SAME', 'VALID', or an explicit tuple/list of per-dimension ints is accepted; any other type (e.g. a list of pairs, a string like 'CAUSAL') hits the final else-branch and raises ValueError 'Invalid padding mode'.

Source

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

  Returns:
    2-tuple: ints: before and after padding for transposed convolution.
  """
  if padding == 'SAME':
    pad_len = k + s - 2
    if s > k - 1:
      pad_a = k - 1
    else:
      pad_a = int(np.ceil(pad_len / 2))
  elif padding == 'VALID':
    pad_len = k + s - 2 + max(k - s, 0)
    pad_a = k - 1
  elif isinstance(padding, tuple):
    pads = tuple(k - p - 1 for p in padding)
    pad_a = pads[0]
    pad_len = sum(pads)
  else:
    raise ValueError(f"Invalid padding mode: {padding}")
  pad_b = pad_len - pad_a
  return pad_a, pad_b

def _flip_axes(x, axes):
  """Flip ndarray 'x' along each axis specified in axes tuple."""
  for axis in axes:
    x = np.flip(x, axis)
  return x


def conv_transpose(lhs: Array, rhs: Array, strides: Sequence[int],
                   padding: str | Sequence[tuple[int, int]],
                   rhs_dilation: Sequence[int] | None = None,
                   dimension_numbers: ConvGeneralDilatedDimensionNumbers = None,
                   transpose_kernel: bool = False,
                   precision: lax.PrecisionLike = None,
                   preferred_element_type: DTypeLike | None = None,
                   use_consistent_padding: bool = False) -> Array:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Pass per-dimension integer padding to conv_transpose: padding=(1, 1)
  2. Pass 'SAME' or 'VALID'
  3. For full control, call conv_general_dilated with explicit (low, high) pairs and manual dilation/strides

Example fix

// before
lax.conv_transpose(x, k, strides=(2,2), padding=[(1,1),(1,1)])
// after
lax.conv_transpose(x, k, strides=(2,2), padding=(1,1))
Defensive patterns

Strategy: validation

Validate before calling

assert padding in ('SAME', 'VALID') or (not isinstance(padding, str) and all(isinstance(p, int) for p in padding)), padding

Prevention

When it happens

Trigger: Calling lax.conv_transpose(..., padding=[(1,1),(1,1)]) with (low,high) pair format, or an unrecognized string; conv_transpose wants per-dim ints like (1,1), unlike conv_general_dilated.

Common situations: Reusing conv_general_dilated-style padding arguments with conv_transpose; copying pair-style padding from flax/haiku layer configs into a direct lax.conv_transpose call.

Related errors


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