jax-ml/jax · error · ValueError

Dimension size after padding is not at least 0, got result s

Error message

Dimension size after padding is not at least 0, got result shape {result}, for padding_config {padding_config} and operand shape {op_shape}

What it means

After applying each axis's (low, high, interior) padding, every output dimension is low + high + dilate(dim, interior+1); if any result is negative (low/high negative enough to over-crop), the pad shape rule raises this ValueError with the computed result shape. lax.pad permits small negative low/high but not so negative that the axis becomes negative.

Source

Thrown at jax/_src/lax/lax.py:7641

def _pad_shape_rule(operand, padding_value, *, padding_config):
  if np.ndim(padding_value) != 0:
    raise ValueError(f"padding_value must be a scalar; got {np.shape(padding_value)=}")
  op_shape = np.shape(operand)
  if not len(padding_config) == np.ndim(operand):
    raise ValueError("length of padding_config must equal the number of axes "
                     f"of operand, got padding_config {padding_config} "
                     f"for operand shape {op_shape}")
  if not all(i >= 0 for _, _, i in padding_config):
    raise ValueError("interior padding in padding_config must be nonnegative, "
                     f"got padding_config {padding_config}")
  result = tuple(l + h + core.dilate_dim(d, i + 1)
                 for (l, h, i), d in zip(padding_config, op_shape))
  if not all(d >= 0 for d in result):
    msg = (f"Dimension size after padding is not at least 0, "
           f"got result shape {result}, for padding_config {padding_config}"
           f" and operand shape {op_shape}")
    raise ValueError(msg)
  return result

def _pad_sharding_rule(operand, padding_value, *, padding_config):
  # TODO(yashkatariya): Once JAX supports uneven sharding at the top level,
  # change this logic to `return operand.sharding` directly.
  out_shape = _pad_shape_rule(operand, padding_value,
                              padding_config=padding_config)
  return slicing._get_sharding_for_varying_out_shape(
      out_shape, operand, 'padding')

def _pad_ur_rule(operand, padding_value, *, padding_config):
  out_unreduced = core.getu(operand)
  kind = UnreducedKind.sum if out_unreduced else None
  return out_unreduced, core.getr(operand), kind

def _pad_transpose(t, operand, padding_value, *, padding_config):
  if type(t) is ad_util.Zero:
    t_operand = ad_util.Zero(operand.aval) if ad.is_undefined_primal(operand) else None

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Clamp negative low/high so the axis stays >= 0, or crop with slicing instead
  2. Compute pads from actual shape: pad = min(pad, x.shape[axis]) and assert result dims >= 0
  3. For symmetric 'same' conv padding use lax.conv_general_dilated's pad handling instead of manual lax.pad

Example fix

# before
y = jax.lax.pad(x, 0, [(-8, -8, 0)])  # dim 4 -> -12
# after
crop = min(8, x.shape[0])
y = x[crop: x.shape[0]-crop] if crop else x
Defensive patterns

Strategy: validation

Validate before calling

result = tuple(l + h + (d * (i + 1) if i else d)
                 for (l, h, i), d in zip(padding_config, x.shape))
assert all(d >= 0 for d in result), result

Type guard

def nonnegative_result(x, config) -> bool:
    return all(l + h + (d if i == 0 else d * (i + 1)) >= 0
               for (l, h, i), d in zip(config, x.shape))

Prevention

When it happens

Trigger: lax.pad(x, 0, [(-10, -10, 0)]) on an axis of length 4 — result dimension -16; computed negative pad widths meant to crop but overshooting.

Common situations: Implementing cropping via negative padding where crop exceeds the size; kernel-size arithmetic producing too-negative 'same' pads; dynamic shapes shrinking under fixed pad widths.

Related errors


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