jax-ml/jax · error · TypeError

`{name}` entries must be the same shape: {nvals}

Error message

`{name}` entries must be the same shape: {nvals}

What it means

jnp.pad converts pad_width/constant_values-like arguments into a concrete numpy array via tree_map. If the entries have inhomogeneous shapes (e.g. a Python list like [[1,2],[3]]), numpy >=1.24 raises 'array has an inhomogeneous shape', and JAX re-raises it as a TypeError explaining the `{name}` entries must all share the same shape.

Source

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

type PadValueLike[T] = Union[T, Sequence[T], Sequence[Sequence[T]]]
type PadValue[T] = tuple[tuple[T, T], ...]

class PadStatFunc(Protocol):
  def __call__(self, array: ArrayLike, /, *,
               axis: int | None = None,
               keepdims: bool = False) -> Array: ...


def _broadcast_to_pairs(nvals: PadValueLike[Any], nd: int, name: str) -> PadValue[Any]:
  try:
    nvals = np.asarray(tree_map(
      lambda x: core.concrete_or_error(None, x, context=f"{name} argument of jnp.pad"),
      nvals))
  except ValueError as e:
    # In numpy 1.24
    if "array has an inhomogeneous shape" in str(e):
      raise TypeError(f'`{name}` entries must be the same shape: {nvals}') from e
    raise

  def as_scalar_dim(v):
    if core.is_dim(v) or not np.shape(v):
      return v
    else:
      raise TypeError(f'`{name}` entries must be the same shape: {nvals}')

  if nvals.shape == (nd, 2):
    # ((before_1, after_1), ..., (before_N, after_N))
    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])

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Make all entries the same shape: use (before, after) pairs for every axis or a single scalar/pair
2. If building pad_width programmatically, normalize each entry to a tuple of two ints before calling jnp.pad
3. Pass a numpy array or use np.broadcast_to on your width specification

Example fix

// before
jnp.pad(x, [[1, 2], [3]])
// after
jnp.pad(x, [[1, 2], [3, 3]])

pad_w = np.array([[1,2],[3,3]])
assert pad_w.dtype != object and len({len(r) for r in pad_w}) == 1, 'ragged pad_width'

def is_regular_pairs(w) -> bool:
    return isinstance(w, (list, tuple)) and all(isinstance(r, (list, tuple)) and len(r) == len(w[0]) for r in w)

When it happens

Trigger: Calling jnp.pad with pad_width or constant_values containing ragged nested sequences, e.g. jnp.pad(x, [[1,2],[3]]) or constant_values=((0,1),(2,)).

Common situations: Dynamically building pad_width from per-axis lists where some axes got a scalar and others got a pair; passing a ragged list constructed in a loop.

Related errors

• multi_dot: last dimension of each array must match first dim
• expected x and y to have same length
• Weights shape must match 'a' shape when axis is None.
• Weights shape {weights.shape} must match reduction axes {tup
• Number of tensordot axes (axes {}) exceeds input ranks ({} a

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