jax-ml/jax · error · ValueError

Cannot apply '{}' padding to empty axis

Error message

Cannot apply '{}' padding to empty axis

What it means

Padding modes 'wrap'/'symmetric'/'reflect'/'edge' are implemented with slicing that requires a nonempty axis; if any padded axis has size 0 and nonzero padding is requested, JAX raises ValueError because there is no data to reflect/wrap.

Source

Thrown at jax/_src/numpy/lax_numpy.py:3932

    v1_2 = as_scalar_dim(nvals[0]), as_scalar_dim(nvals[1])
    return tuple(v1_2 for i in range(nd))
  elif nvals.shape == (1,):
    # (pad,)
    v = as_scalar_dim(nvals[0])
    return tuple((v, v) for i in range(nd))
  elif nvals.shape == ():
    # pad
    v = as_scalar_dim(nvals.flat[0])
    return tuple((v, v) for i in range(nd))
  else:
    raise ValueError(f"jnp.pad: {name} with {nd=} has unsupported shape {nvals.shape}. "
                     f"Valid shapes are ({nd}, 2), (1, 2), (2,), (1,), or ().")


def _check_no_padding(axis_padding: tuple[Any, Any], mode: str):
  if (axis_padding[0] > 0 or axis_padding[1] > 0):
    msg = "Cannot apply '{}' padding to empty axis"
    raise ValueError(msg.format(mode))


def _pad_constant(array: Array, pad_width: PadValue[int], constant_values: Array) -> Array:
  nd = np.ndim(array)
  constant_values = lax._convert_element_type(
      constant_values, array.dtype, dtypes.is_weakly_typed(array))
  constant_values_nd = np.ndim(constant_values)

  if constant_values_nd == 0:
    widths = [(low, high, 0) for (low, high) in pad_width]
    return lax.pad(array, constant_values, widths)

  if constant_values_nd == 1:
    if constant_values.shape[-1] == 1:
      widths = [(low, high, 0) for (low, high) in pad_width]
      return lax.pad(array, squeeze(constant_values), widths)
    elif constant_values.shape[-1] != 2:
      raise ValueError("jnp.pad: constant_values has unsupported shape "

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Use mode='constant' when the array may have empty axes
  2. Guard the empty case: skip padding or pad only nonempty axes
  3. Fix upstream logic that produces zero-length axes

Example fix

// before
jnp.pad(x, ((2, 2), (0, 0)), mode='wrap')
// after
mode = 'wrap' if x.shape[0] > 0 else 'constant'
jnp.pad(x, ((2, 2), (0, 0)), mode=mode)
Defensive patterns

Strategy: type-guard

Validate before calling

mode = 'wrap'
if any(s == 0 and (b > 0 or a > 0) for s, (b, a) in zip(x.shape, pad_width)) and mode != 'constant':
    mode = 'constant'

Type guard

def can_nonconstant_pad(x, pad_width) -> bool:
    return all(sz > 0 or (b == 0 and a == 0) for sz, (b, a) in zip(x.shape, pad_width))

Prevention

When it happens

Trigger: jnp.pad(jnp.zeros((0, 3)), ((2, 2), (0, 0)), mode='wrap') or mode='symmetric'/'reflect'/'edge' on an empty axis with padding > 0.

Common situations: Batch pipelines where a dynamic batch or sequence dimension can collapse to zero; filtering/selection producing empty arrays before padding.

Related errors


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