{"record":{"id":"e0ceeee36f4d5e64","repo":"jax-ml/jax","slug":"clamp-requires-max-shape-operand-shape-or-max-s","errorCode":null,"errorMessage":"clamp requires max.shape == operand.shape or max.shape == (), got max.shape={max.shape}, {operand.shape=}.","messagePattern":"clamp requires max\\.shape == operand\\.shape or max\\.shape == \\(\\), got max\\.shape=(.+?), (.+?)\\.","errorType":"validation","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/lax.py","lineNumber":7165,"sourceCode":"  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\n  # avoid transposes and some broadcasts in special cases\n  if min_bdim == x_bdim == max_bdim:\n    if np.shape(min) == np.shape(x) == np.shape(max):","sourceCodeStart":7147,"sourceCodeEnd":7183,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/lax.py#L7147-L7183","documentation":"lax.clamp's max bound must be scalar (shape ()) or match the operand shape exactly. Non-scalar, mismatched max bounds are rejected.","triggerScenarios":"Calling jax.lax.clamp(min, x, max) where max.shape is neither () nor equal to x.shape.","commonSituations":"Passing per-feature upper bounds (e.g. (d,) tensor) to a (b, d) operand; assuming numpy-style broadcasting semantics in clamp.","solutions":["Broadcast max to the operand shape before the call","Use a 0-d scalar for a global bound: lax.clamp(lo, x, jnp.asarray(hi))","Or replace with broadcast-friendly jnp.minimum/jnp.maximum composition"],"exampleFix":"// before\nout = lax.clamp(0.0, x, jnp.full((16,), 1.0))   # x is (8,16)\n// after\nhi = jnp.broadcast_to(1.0, x.shape)\nout = lax.clamp(0.0, x, hi)","handlingStrategy":"validation","validationCode":"if max.shape and max.shape != x.shape:\n    max = jnp.broadcast_to(max, x.shape)","typeGuard":"def clamp_bound_ok(b, x) -> bool:\n    return b.shape == () or b.shape == x.shape","tryCatchPattern":null,"preventionTips":["Scalarize global bounds: jnp.asarray(value)","Broadcast per-feature bounds to full shape before clamp"],"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"}