jax-ml/jax · error · ValueError

Expected pad_width to have shape {(nd, 2)}; got {pad_width_a

Error message

Expected pad_width to have shape {(nd, 2)}; got {pad_width_arr.shape}.

What it means

After _broadcast_to_pairs normalization, jnp.pad requires the resulting pad_width array to have exactly shape (nd, 2) (one (before, after) per dimension). This check catches width specifications that normalized to a different shape.

Source

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

         constant_values: ArrayLike, stat_length: PadValueLike[int] | None,
         end_values: PadValueLike[ArrayLike], reflect_type: str):
  array = asarray(array)
  nd = np.ndim(array)

  if nd == 0:
    return array

  stat_funcs: dict[str, PadStatFunc] = {
      "maximum": reductions.amax,
      "minimum": reductions.amin,
      "mean": reductions.mean,
      "median": reductions.median
  }

  pad_width = _broadcast_to_pairs(pad_width, nd, "pad_width")
  pad_width_arr = np.array(pad_width)
  if pad_width_arr.shape != (nd, 2):
    raise ValueError(f"Expected pad_width to have shape {(nd, 2)}; got {pad_width_arr.shape}.")

  if np.any(pad_width_arr < 0):
    raise ValueError("index can't contain negative values")

  if mode == "constant":
    return _pad_constant(array, pad_width, asarray(constant_values))

  elif mode == "wrap":
    return _pad_wrap(array, pad_width)

  elif mode in ("symmetric", "reflect"):
    return _pad_symmetric_or_reflect(array, pad_width, str(mode), reflect_type)

  elif mode == "edge":
    return _pad_edge(array, pad_width)

  elif mode == "linear_ramp":
    end_values = _broadcast_to_pairs(end_values, nd, "end_values")

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Convert pad_width to a plain (nd, 2) int array or list of pairs before calling jnp.pad
  2. Avoid mixing arrays and scalars inside pad_width

Example fix

// before
jnp.pad(x, np.array([[1, 2], [np.array(3), 4]], dtype=object))
// after
jnp.pad(x, np.array([[1, 2], [3, 4]]))
Defensive patterns

Strategy: validation

Validate before calling

w = np.array(pad_width)
assert w.shape == (x.ndim, 2), f'pad_width must be {(x.ndim, 2)}, got {w.shape}'

Prevention

When it happens

Trigger: Rare: pad_width whose conversion via np.array yields unexpected shape, e.g. mixed types producing an object array or extra dimensions beyond the handled shapes in _broadcast_to_pairs.

Common situations: Passing pad_width as an object array of mixed scalars/arrays; inconsistent nesting depth.

Related errors


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