{"record":{"id":"2952d1ec4b733e6b","repo":"jax-ml/jax","slug":"broadcast-in-dim-broadcast-dimensions-must-be-a-su","errorCode":null,"errorMessage":"broadcast_in_dim broadcast_dimensions must be a subset of output dimensions, got {} for operand ndim {} and shape {}.","messagePattern":"broadcast_in_dim broadcast_dimensions must be a subset of output dimensions, got (.+?) for operand ndim (.+?) and shape (.+?)\\.","errorType":"validation","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/lax.py","lineNumber":6925,"sourceCode":"\ndef _broadcast_in_dim_shape_rule(operand, *, shape, broadcast_dimensions,\n                                 sharding):\n  _check_shapelike('broadcast_in_dim', 'shape', shape)\n  _check_shapelike('broadcast_in_dim', 'broadcast_dimensions',\n                   broadcast_dimensions)\n  operand_ndim = np.ndim(operand)\n  if operand_ndim != len(broadcast_dimensions):\n    msg = ('broadcast_in_dim broadcast_dimensions must have length equal to '\n           'operand ndim; got broadcast_dimensions {} for operand ndim {}.')\n    raise TypeError(msg.format(broadcast_dimensions, operand_ndim))\n  if len(shape) < operand_ndim:\n    msg = ('broadcast_in_dim target broadcast shape must have equal or higher rank '\n           'to the operand shape; got operand ndim {} and target broadcast ndim {}.')\n    raise TypeError(msg.format(operand_ndim, len(shape)))\n  if not set(broadcast_dimensions).issubset(set(range(len(shape)))):\n    msg = ('broadcast_in_dim broadcast_dimensions must be a subset of output '\n           'dimensions, got {} for operand ndim {} and shape {}.')\n    raise TypeError(msg.format(broadcast_dimensions, operand_ndim, shape))\n  if not all(core.definitely_equal_one_of_dim(operand.shape[i],\n                                              [1, shape[broadcast_dimensions[i]]])\n             for i in range(operand_ndim)):\n    msg = (\n        \"broadcast_in_dim operand dimension sizes must either be 1, or be \"\n        \"equal to their corresponding dimensions in the target broadcast \"\n        \"shape; got operand of shape {}, target broadcast shape {}, \"\n        \"broadcast_dimensions {} \")\n    raise TypeError(msg.format(\n        tuple(core.replace_tracer_for_error_message(d) for d in operand.shape),\n        shape, broadcast_dimensions))\n  if len(broadcast_dimensions) != len(set(broadcast_dimensions)):\n    msg = (\"broadcast_in_dim broadcast_dimensions must not contain duplicates, \"\n           \"got broadcast_dimensions {}\")\n    raise TypeError(msg.format(broadcast_dimensions))\n  return shape\n\ndef _broadcast_in_dim_sharding_rule(operand, *, shape, broadcast_dimensions,","sourceCodeStart":6907,"sourceCodeEnd":6943,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/lax.py#L6907-L6943","documentation":"Every index in broadcast_dimensions must lie within the output shape's dimensions (0..len(shape)-1). Indices outside that range (negative or too large) are invalid because they reference non-existent output dims.","triggerScenarios":"Calling broadcast_in_dim(operand, shape, broadcast_dimensions) where some index >= len(shape) or < 0 (e.g. index 3 for a 3-D output).","commonSituations":"Off-by-one when mapping dims for a rank-(n+1) output (using indices 1..n+1 instead of 0..n); copying broadcast_dimensions from a differently-shaped call site.","solutions":["Validate indices: all(0 <= i < len(shape) for i in broadcast_dimensions)","Remember indices are into the OUTPUT shape, and the max valid index is len(shape)-1","Write a small helper that derives broadcast_dimensions from operand/output shapes"],"exampleFix":"// before\nx = jnp.zeros((3,))\ny = lax.broadcast_in_dim(x, (2, 3, 4), (1, 3))  # 3 out of range for rank-3\n// after\ny = lax.broadcast_in_dim(x, (2, 3, 4), (1, 2))","handlingStrategy":"validation","validationCode":"assert all(0 <= i < len(shape) for i in broadcast_dimensions)","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Indices are into the OUTPUT shape; max is len(shape)-1"],"tags":["jax","broadcast-in-dim","index-out-of-range"],"backgroundTag":"invalid-broadcast-dimensions","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}