jax-ml/jax · error · TypeError

pad operand and padding_value must be same dtype: got {} and

Error message

pad operand and padding_value must be same dtype: got {} and {}.

What it means

jax.lax.pad (used by jnp.pad's 'constant'-like internal path and directly) requires padding_value to have exactly the same dtype as the operand, because the output dtype is the operand dtype and the fill value must be representable in it. A mismatch raises this TypeError at the dtype-rule stage.

Source

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

  out_vma = core.standard_vma_rule('split', operand)
  out_shapes = _split_shape_rule(operand, sizes=sizes, axis=axis)
  return [out_vma] * len(out_shapes)

split_p = core.Primitive('split')
split_p.multiple_results = True
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, "

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Cast the padding value: jax.lax.pad(x, jnp.asarray(0, x.dtype), ...);
2. Prefer jnp.pad(x, pads, mode='constant', constant_values=...) which handles promotion
3. Standardize dtypes at model input boundaries to avoid mixed-dtype constants

Example fix

# before
y = jax.lax.pad(x_bf16, 0.0, config)
# after
y = jax.lax.pad(x_bf16, jnp.asarray(0.0, x_bf16.dtype), config)

pad_val = jnp.asarray(pad_val, x.dtype)  # or jnp.result_type check

def same_dtype(x, v) -> bool:
    return jnp.result_type(x) == jnp.result_type(v)

When it happens

Trigger: jax.lax.pad(jnp.zeros(3, jnp.float32), 0) — Python int weakly typed to int32/weak but still mismatched under some promotion contexts; padding a float32 array with jnp.int32(0); padding a bfloat16 array with a float32 scalar.

Common situations: Padding with default Python 0 or 1 on half-precision (bfloat16) tensors; mixing dtypes after enabling legacy or strict promotion settings; padding int arrays with float values.

Related errors

• Incorrect output dtype for return value #{i}: Expected: {out
• {} function carry input and carry output must have equal typ
• lower and upper arguments to fori_loop must have equal types
• the first argument to pow must have an inexact dtype (float
• {name} argument type error: {lhs.dtype}, {rhs.dtype}

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