jax-ml/jax · error · ValueError

padding_value must be a scalar; got {np.shape(padding_value)

Error message

padding_value must be a scalar; got {np.shape(padding_value)=}

What it means

The padding_value in jax.lax.pad must be a scalar (0-D). If an array with any non-empty shape is passed, the pad shape rule raises this ValueError echoing np.shape(padding_value). The value is broadcast into every padded cell, so non-scalars are meaningless here.

Source

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

split_p.def_abstract_eval(
    partial(standard_multi_result_abstract_eval, split_p, _split_shape_rule,
            _split_dtype_rule, _split_weak_type_rule, _split_sharding_rule,
            _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):

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Use a true scalar: jax.lax.pad(x, jnp.asarray(0, x.dtype), config) or the Python scalar 0
  2. Put per-axis low/high/interior amounts in padding_config, not in padding_value
  3. For constant-valued padding use jnp.pad(x, widths, constant_values=c)

Example fix

# before
y = jax.lax.pad(x, jnp.zeros((1,)), config)
# after
y = jax.lax.pad(x, 0, config)
Defensive patterns

Strategy: validation

Validate before calling

assert np.ndim(padding_value) == 0, np.shape(padding_value)

Type guard

def is_scalar(v) -> bool:
    return np.ndim(v) == 0

Prevention

When it happens

Trigger: jax.lax.pad(x, jnp.zeros((1,)), config); passing a length-1 vector or a per-axis padding array as padding_value.

Common situations: Confusing per-axis pad widths (which go in padding_config) with the fill value; reusing a weights array as fill; passing constant_values arrays from another API.

Related errors


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