jax-ml/jax · error · TypeError

clamp requires min.shape == operand.shape or min.shape == ()

Error message

clamp requires min.shape == operand.shape or min.shape == (), got min.shape={min.shape}, {operand.shape=}.

What it means

lax.clamp's min bound must either be a scalar (shape ()) or exactly the same shape as the operand. Any other shape is rejected because per-element clamping requires aligned bounds.

Source

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

def _tile_batch_rule(batched_args, batch_dims, *, reps):
  operand, = batched_args
  bdim, = batch_dims
  new_reps = list(reps)
  new_reps.insert(bdim, 1)
  return tile(operand, reps=new_reps), bdim

tile_p = core.Primitive('tile')
tile_p.def_abstract_eval(_tile_abstract_eval)
tile_p.def_impl(partial(dispatch.apply_primitive, tile_p))
ad.deflinear2(tile_p, _tile_transpose_rule)
batching.primitive_batchers[tile_p] = _tile_batch_rule
mlir.register_lowering(tile_p, _tile_lower)


def _clamp_shape_rule(min, operand, max):
  if min.shape and min.shape != operand.shape:
    raise TypeError("clamp requires min.shape == operand.shape or min.shape == "
                    f"(), got min.shape={min.shape}, {operand.shape=}.")
  if max.shape and max.shape != operand.shape:
    raise TypeError("clamp requires max.shape == operand.shape or max.shape == "
                    f"(), got max.shape={max.shape}, {operand.shape=}.")
  return operand.shape

def _clamp_sharding_rule(min, operand, max):
  return operand.sharding

_clamp_dtype_rule = partial(naryop_dtype_rule, input_dtype, [_any, _any, _any],
                            'clamp')

def _clamp_batch_rule(batched_args, batch_dims, **params):
  min, x, max = batched_args
  min_bdim, x_bdim, max_bdim = batch_dims
  size = next(x.shape[i] for x, i in zip(batched_args, batch_dims)
              if i is not None)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Broadcast min explicitly first: min = jnp.broadcast_to(min, x.shape), or use lax.clamp(min_b, x, max_b) after broadcasting
2. If min is meant to be global, pass a 0-d scalar: jnp.asarray(lo) or float constant
3. Alternatively use jnp.minimum(jnp.maximum(x, min), max) which does broadcast

Example fix

// before
x = jnp.zeros((8, 16))
out = lax.clamp(jnp.zeros(16), x, jnp.ones(16))  # min not scalar, not full shape
// after
lo = jnp.broadcast_to(jnp.zeros(16), x.shape)
hi = jnp.broadcast_to(jnp.ones(16), x.shape)
out = lax.clamp(lo, x, hi)

if min.shape and min.shape != x.shape:
    min = jnp.broadcast_to(min, x.shape)

def clamp_bound_ok(b, x) -> bool:
    return b.shape == () or b.shape == x.shape

When it happens

Trigger: Calling jax.lax.clamp(min, x, max) where min is non-scalar and min.shape != x.shape, e.g. per-row minima with shape (d,) against x of shape (b, d) without broadcasting.

Common situations: Expecting clamp to broadcast like jnp.minimum; passing 1-D bounds to a 2-D tensor (common in attention masking or value clipping).

Related errors

• clamp requires max.shape == operand.shape or max.shape == ()
• put_along_axis argument 'values' must be broadcastable to 'i
• x and y cannot be broadcast together.
• Cannot broadcast shape {self.shape} to layout {o.layout}
• unexpected JAX type (e.g. shape/dtype) for argument to VJP f

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