jax-ml/jax · error · TypeError

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

Error message

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

What it means

lax.clamp's max bound must be scalar (shape ()) or match the operand shape exactly. Non-scalar, mismatched max bounds are rejected.

Source

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

  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)

  # avoid transposes and some broadcasts in special cases
  if min_bdim == x_bdim == max_bdim:
    if np.shape(min) == np.shape(x) == np.shape(max):

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Broadcast max to the operand shape before the call
2. Use a 0-d scalar for a global bound: lax.clamp(lo, x, jnp.asarray(hi))
3. Or replace with broadcast-friendly jnp.minimum/jnp.maximum composition

Example fix

// before
out = lax.clamp(0.0, x, jnp.full((16,), 1.0))   # x is (8,16)
// after
hi = jnp.broadcast_to(1.0, x.shape)
out = lax.clamp(0.0, x, hi)

if max.shape and max.shape != x.shape:
    max = jnp.broadcast_to(max, 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 max.shape is neither () nor equal to x.shape.

Common situations: Passing per-feature upper bounds (e.g. (d,) tensor) to a (b, d) operand; assuming numpy-style broadcasting semantics in clamp.

Related errors

• clamp requires min.shape == operand.shape or min.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/e0ceeee36f4d5e64. Report an issue: GitHub.