jax-ml/jax · error · ValueError

String padding is not implemented for transposed convolution

Error message

String padding is not implemented for transposed convolution using this op. Please either exactly specify the required padding or use conv_transpose.

What it means

conv_general_dilated cannot compute automatic padding ('SAME'/'VALID') for a transposed convolution, which is what lhs_dilation != all-ones implies. If string padding is combined with non-unit lhs_dilation the op raises ValueError and directs you to conv_transpose, which handles padding computation for transposed convolutions.

Source

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

  For example, to indicate dimension numbers consistent with the ``conv``
  function with two spatial dimensions, one could use ``('NCHW', 'OIHW',
  'NCHW')``. As another example, to indicate dimension numbers consistent with
  the TensorFlow Conv2D operation, one could use ``('NHWC', 'HWIO', 'NHWC')``.
  When using the latter form of convolution dimension specification, window
  strides are associated with spatial dimension character labels according to
  the order in which the labels appear in the ``rhs_spec`` string, so that
  ``window_strides[0]`` is matched with the dimension corresponding to the first
  character appearing in rhs_spec that is not ``'I'`` or ``'O'``.

  If ``dimension_numbers`` is ``None``, the default is ``('NCHW', 'OIHW',
  'NCHW')`` (for a 2D convolution).
  """
  dnums = conv_dimension_numbers(lhs.shape, rhs.shape, dimension_numbers)
  out_sharding = canonicalize_sharding(out_sharding, 'dot_general')
  if lhs_dilation is None:
    lhs_dilation = (1,) * (lhs.ndim - 2)
  elif isinstance(padding, str) and not len(lhs_dilation) == lhs_dilation.count(1):
    raise ValueError(
        "String padding is not implemented for transposed convolution "
        "using this op. Please either exactly specify the required padding or "
        "use conv_transpose.")
  if rhs_dilation is None:
    rhs_dilation = (1,) * (rhs.ndim - 2)
  if isinstance(padding, str):
    lhs_perm, rhs_perm, _ = dnums
    rhs_shape = np.take(rhs.shape, rhs_perm)[2:]
    effective_rhs_shape = [core.dilate_dim(k, r) for k, r in zip(rhs_shape, rhs_dilation)]
    padding = lax.padtype_to_pads(
        np.take(lhs.shape, lhs_perm)[2:], effective_rhs_shape,
        window_strides, padding)
  else:
    try:
      padding = tuple((operator.index(lo), operator.index(hi))
                      for lo, hi in padding)
    except (ValueError, TypeError) as e:
      raise ValueError(

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Use lax.conv_transpose, which computes correct transposed-conv padding for 'SAME'/'VALID'
  2. Specify explicit padding as a list of (low, high) pairs instead of a string

Example fix

// before
out = lax.conv_general_dilated(x, k, strides=(2,2), padding='SAME', lhs_dilation=(2,2))
// after
out = lax.conv_transpose(x, k, strides=(2,2), padding='SAME')
Defensive patterns

Strategy: fallback

Validate before calling

if isinstance(padding, str) and any(d != 1 for d in (lhs_dilation or (1,)*(lhs.ndim-2))):
    use_conv_transpose = True  # switch API

Prevention

When it happens

Trigger: Calling lax.conv_general_dilated(lhs, rhs, strides, 'SAME', lhs_dilation=(2,2)) directly; indirectly by calling lax.conv_transpose with a non-default padding is fine, but hand-rolled transposed convs through conv_general_dilated with string padding hit this.

Common situations: Manually implementing transposed conv / upsampling via dilation instead of using lax.conv_transpose; porting PyTorch ConvTranspose2d semantics to JAX by hand.

Related errors


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