jax-ml/jax · error · ValueError

length of padding_config must equal the number of axes of op

Error message

length of padding_config must equal the number of axes of operand, got padding_config {padding_config} for operand shape {op_shape}

What it means

padding_config for jax.lax.pad must contain exactly one (low, high, interior) triple per operand dimension. This ValueError fires when len(padding_config) != operand.ndim, e.g. 2 triples for a 3-D array.

Source

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

            _split_vma_rule, _split_ur_rule, None))
split_p.def_impl(partial(dispatch.apply_primitive, split_p))
ad.deflinear2(split_p, _split_transpose_rule)
batching.primitive_batchers[split_p] = _split_batch_rule
mlir.register_lowering(split_p, _split_lower)

def _pad_dtype_rule(operand, padding_value, *, padding_config):
  if operand.dtype != padding_value.dtype:
    msg = "pad operand and padding_value must be same dtype: got {} and {}."
    raise TypeError(msg.format(operand.dtype, padding_value.dtype))

  return input_dtype(operand, padding_value)

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,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Build the config from the operand's ndim: config = tuple(pad for _ in range(x.ndim))
2. Derive per-axis widths programmatically instead of hardcoding
3. Check x.ndim and len(padding_config) with an assert before calling lax.pad

Example fix

# before
config = [(1,1,0), (1,1,0)]
y = jax.lax.pad(x3d, 0, config)  # x3d.ndim == 3
# after
config = [(1,1,0)] * x3d.ndim
y = jax.lax.pad(x3d, 0, config)

assert len(padding_config) == x.ndim, (len(padding_config), x.ndim)

def config_matches_rank(config, x) -> bool:
    return len(config) == x.ndim

When it happens

Trigger: Passing [(1,1,0),(1,1,0)] for a rank-3 tensor; reusing a 2-D conv padding config on 3-D input; hardcoding pad widths after adding a batch/channel axis.

Common situations: Adding channels-first/batch dims to a pipeline without updating pad configs; adapting image padding code for video (rank+1); configs shared across models with different ranks.

Related errors

• unstack requires arrays with rank > 0, however a scalar arra
• pad operand and padding_value must be same dtype: got {} and
• padding_value must be a scalar; got {np.shape(padding_value)
• interior padding in padding_config must be nonnegative, got
• Dimension size after padding is not at least 0, got result s

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