{"record":{"id":"4640ced5fd30d6bd","repo":"jax-ml/jax","slug":"broadcast-in-dim-operand-dimension-sizes-must-eith","errorCode":null,"errorMessage":"broadcast_in_dim operand dimension sizes must either be 1, or be equal to their corresponding dimensions in the target broadcast shape; got operand of shape {}, target broadcast shape {}, broadcast_dimensions {} ","messagePattern":"broadcast_in_dim operand dimension sizes must either be 1, or be equal to their corresponding dimensions in the target broadcast shape; got operand of shape (.+?), target broadcast shape (.+?), broadcast_dimensions (.+?) ","errorType":"validation","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/lax.py","lineNumber":6934,"sourceCode":"           '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,\n                                    sharding):\n  if sharding is not None:\n    return sharding\n  bds = set(broadcast_dimensions)\n  orig_spec = iter(operand.sharding.spec.partitions)\n  new_spec = [next(orig_spec) if i in bds else None for i in range(len(shape))]\n  assert next(orig_spec, None) is None\n  mesh = (get_abstract_mesh() if operand.sharding.mesh.empty else\n          operand.sharding.mesh)","sourceCodeStart":6916,"sourceCodeEnd":6952,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/lax.py#L6916-L6952","documentation":"Each operand dimension must map to an output dimension of the same size, unless the operand dim is 1 (which broadcasts). If an operand dim size matches neither 1 nor the target dim size, the broadcast is invalid.","triggerScenarios":"Calling broadcast_in_dim where operand.shape[i] != shape[broadcast_dimensions[i]] and operand.shape[i] != 1.","commonSituations":"Assuming broadcast_in_dim tiles/repeats a non-unit dim (it cannot — use lax.tile for that); shape typos where the target dim differs from the operand dim by a constant.","solutions":["Fix the target shape so dim sizes match the operand, or make the operand dim 1 before broadcasting","Use lax.tile / jnp.tile for repetition of non-unit dims","Check the mapping index: often the wrong output index pairs the operand dim with an unrelated dim"],"exampleFix":"// before\nx = jnp.zeros((4, 3))\ny = lax.broadcast_in_dim(x, (4, 5), (0, 1))  # operand dim 3 vs target 5\n// after\nx1 = x[:, None]                              # (4, 1)\ny = lax.broadcast_in_dim(x1, (4, 5), (0, 1))  # 1 broadcasts to 5","handlingStrategy":"validation","validationCode":"ok = all(operand.shape[i] in (1, shape[bd[i]]) for i in range(np.ndim(operand)))\nassert ok","typeGuard":"def can_broadcast(operand_shape, shape, bd) -> bool:\n    return all(operand_shape[i] == 1 or operand_shape[i] == shape[bd[i]]\n               for i in range(len(operand_shape)))","tryCatchPattern":null,"preventionTips":["Use lax.tile for repetition, not broadcast_in_dim","Insert size-1 dims via reshape/expand_dims before broadcasting"],"tags":["jax","broadcast-in-dim","shape-mismatch"],"backgroundTag":"shape-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}