jax-ml/jax · error · ValueError

jnp.pad: {name} with {nd=} has unsupported shape {nvals.shap

Error message

jnp.pad: {name} with {nd=} has unsupported shape {nvals.shape}. Valid shapes are ({nd}, 2), (1, 2), (2,), (1,), or ().

What it means

jnp.pad normalizes pad_width/constant_values/stat_length-like arguments to a shape of (nd, 2). Any input whose numpy shape is not (nd,2), (1,2), (2,), (1,), or () is rejected with this ValueError listing valid shapes.

Source

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

    return tuple((as_scalar_dim(nval[0]), as_scalar_dim(nval[1])) for nval in nvals)
  elif nvals.shape == (1, 2):
    # ((before, after),)
    v1_2 = as_scalar_dim(nvals[0, 0]), as_scalar_dim(nvals[0, 1])
    return tuple(v1_2 for i in range(nd))
  elif nvals.shape == (2,):
    # (before, after)  (not in the numpy docstring but works anyway)
    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)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Provide exactly (before, after) pairs for each of the nd axes, or a single pair/scalar
  2. Check x.ndim and reshape your width array to (x.ndim, 2)

Example fix

// before
jnp.pad(x, (1, 2, 3))
// after
jnp.pad(x, [(1, 2), (1, 2)])  # for 2-D x
Defensive patterns

Strategy: validation

Validate before calling

import numpy as np
w = np.asarray(pad_width)
nd = x.ndim
assert w.shape in [(nd,2),(1,2),(2,),(1,),()], f'bad pad_width shape {w.shape}'

Prevention

When it happens

Trigger: Passing e.g. jnp.pad(x, (1,2,3)) (shape (3,)), or a (4,2) width array for a 2-D input, or a shape like (2,1).

Common situations: Assuming jnp.pad accepts arbitrary per-axis triples; passing more entries than the array has dimensions; transposed width arrays.

Related errors


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