jax-ml/jax · error · ValueError

Shape polymorphism is supported for jnp.pad with 'reflect' o

Error message

Shape polymorphism is supported for jnp.pad with 'reflect' or 'symmetric' padding mode only when it is possible to determine at lowering time that the axis size (= {axis_size}) is larger than 1 and larger or equal than the padding length (= {padding}). Error while handling {'left' if before else 'right'} padding on axis {i}.

What it means

With shape polymorphism (symbolic dimensions), jnp.pad's reflect/symmetric implementation must statically decide that the axis size exceeds 1 and is at least the padding length so the padding loop terminates in one iteration. If the symbolic comparison is inconclusive, this ValueError is raised from the explicit check.

Source

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

      if before:
        edge = lax_slicing.slice_in_dim(array, 0, 1, axis=i)
      else:
        edge = lax_slicing.slice_in_dim(array, -1, None, axis=i)

      # Try to give nicer error messages for unsupported shape polymorphic uses
      shape_poly_error_msg = lambda: (
          "Shape polymorphism is supported for jnp.pad with 'reflect' or "
          "'symmetric' padding mode only when it is possible to determine "
          f"at lowering time that the axis size (= {axis_size}) is larger than 1 "
          f"and larger or equal than the padding length (= {padding}). "
          f"Error while handling {'left' if before else 'right'} padding on axis {i}.")
      try:
        # We check that we can determine all comparisons.
        offset = 1 if (mode == "reflect" and axis_size > 1) else 0
        has_poly_dim = not core.is_constant_shape((axis_size, padding))
        # For shape polymorphism, ensure the loop below ends after 1 iteration
        if has_poly_dim and not (axis_size > 1 and axis_size - offset >= padding):
          raise ValueError(shape_poly_error_msg())
      except core.InconclusiveDimensionOperation as e:
        raise ValueError(shape_poly_error_msg()) from e

      while padding > 0:
        curr_pad = min(padding, axis_size - offset)
        padding -= curr_pad
        if has_poly_dim: assert padding == 0

        if before:
          start = offset
          stop = offset + curr_pad
        else:
          start = -(curr_pad + offset)
          stop = None if (mode == "symmetric" or axis_size == 1) else -1

        x = lax_slicing.slice_in_dim(array, start, stop, axis=i)
        x = flip(x, axis=i)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Reduce padding to at most axis_size - 1 (symmetric) / axis_size - 2 (reflect) as a static constant
2. Use mode='constant' for polymorphic shapes
3. Constrain the polymorphic dimension with explicit bounds so comparisons are conclusive (e.g. use jax.ShapeDtypeStruct with bounded polys)

Example fix

// before
jnp.pad(x, ((n - 1, 0),), mode='symmetric')  # n symbolic
// after
jnp.pad(x, ((2, 0),), mode='symmetric')  # static, small padding

try:
    _ = bool(axis_size > 1 and axis_size - 1 >= padding)
except Exception:
    mode = 'constant'  # cannot verify statically

try:
    y = jnp.pad(x, w, mode='reflect')
except ValueError:
    y = jnp.pad(x, w, mode='constant')

When it happens

Trigger: Using jax.export/jit with symbolic axis sizes (e.g. batch='n') and calling jnp.pad(..., mode='reflect'|'symmetric') with padding >= the symbolic axis size, or where axis_size > 1 cannot be proven.

Common situations: Exporting models with dynamic batch/sequence dimensions that pad with reflection; padding lengths derived from user config larger than the traced dimension.

Related errors

• Cannot apply '{}' padding to empty axis
• Encountered dimension variable '{self.var}' that is not appe
• Cannot divide {self} by {divisor}.
• __pow__ modulo not implemented
• Symbolic dimension cannot be raised to non-integer powers: '

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