{"record":{"id":"a9a119d584abb21a","repo":"jax-ml/jax","slug":"clamp-requires-min-shape-operand-shape-or-min-s","errorCode":null,"errorMessage":"clamp requires min.shape == operand.shape or min.shape == (), got min.shape={min.shape}, {operand.shape=}.","messagePattern":"clamp requires min\\.shape == operand\\.shape or min\\.shape == \\(\\), got min\\.shape=(.+?), (.+?)\\.","errorType":"validation","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/lax.py","lineNumber":7162,"sourceCode":"\ndef _tile_batch_rule(batched_args, batch_dims, *, reps):\n  operand, = batched_args\n  bdim, = batch_dims\n  new_reps = list(reps)\n  new_reps.insert(bdim, 1)\n  return tile(operand, reps=new_reps), bdim\n\ntile_p = core.Primitive('tile')\ntile_p.def_abstract_eval(_tile_abstract_eval)\ntile_p.def_impl(partial(dispatch.apply_primitive, tile_p))\nad.deflinear2(tile_p, _tile_transpose_rule)\nbatching.primitive_batchers[tile_p] = _tile_batch_rule\nmlir.register_lowering(tile_p, _tile_lower)\n\n\ndef _clamp_shape_rule(min, operand, max):\n  if min.shape and min.shape != operand.shape:\n    raise TypeError(\"clamp requires min.shape == operand.shape or min.shape == \"\n                    f\"(), got min.shape={min.shape}, {operand.shape=}.\")\n  if max.shape and max.shape != operand.shape:\n    raise TypeError(\"clamp requires max.shape == operand.shape or max.shape == \"\n                    f\"(), got max.shape={max.shape}, {operand.shape=}.\")\n  return operand.shape\n\ndef _clamp_sharding_rule(min, operand, max):\n  return operand.sharding\n\n_clamp_dtype_rule = partial(naryop_dtype_rule, input_dtype, [_any, _any, _any],\n                            'clamp')\n\ndef _clamp_batch_rule(batched_args, batch_dims, **params):\n  min, x, max = batched_args\n  min_bdim, x_bdim, max_bdim = batch_dims\n  size = next(x.shape[i] for x, i in zip(batched_args, batch_dims)\n              if i is not None)\n","sourceCodeStart":7144,"sourceCodeEnd":7180,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/lax.py#L7144-L7180","documentation":"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.","triggerScenarios":"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.","commonSituations":"Expecting clamp to broadcast like jnp.minimum; passing 1-D bounds to a 2-D tensor (common in attention masking or value clipping).","solutions":["Broadcast min explicitly first: min = jnp.broadcast_to(min, x.shape), or use lax.clamp(min_b, x, max_b) after broadcasting","If min is meant to be global, pass a 0-d scalar: jnp.asarray(lo) or float constant","Alternatively use jnp.minimum(jnp.maximum(x, min), max) which does broadcast"],"exampleFix":"// before\nx = jnp.zeros((8, 16))\nout = lax.clamp(jnp.zeros(16), x, jnp.ones(16))  # min not scalar, not full shape\n// after\nlo = jnp.broadcast_to(jnp.zeros(16), x.shape)\nhi = jnp.broadcast_to(jnp.ones(16), x.shape)\nout = lax.clamp(lo, x, hi)","handlingStrategy":"validation","validationCode":"if min.shape and min.shape != x.shape:\n    min = jnp.broadcast_to(min, x.shape)","typeGuard":"def clamp_bound_ok(b, x) -> bool:\n    return b.shape == () or b.shape == x.shape","tryCatchPattern":null,"preventionTips":["Remember lax.clamp does NOT broadcast","Use jnp.minimum/jnp.maximum when you need broadcasting"],"tags":["jax","clamp","shape-mismatch","broadcasting"],"backgroundTag":"shape-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}