{"record":{"id":"a37ad947274a0b29","repo":"jax-ml/jax","slug":"broadcast-in-dim-broadcast-dimensions-must-have-le","errorCode":null,"errorMessage":"broadcast_in_dim broadcast_dimensions must have length equal to operand ndim; got broadcast_dimensions {} for operand ndim {}.","messagePattern":"broadcast_in_dim broadcast_dimensions must have length equal to operand ndim; got broadcast_dimensions (.+?) for operand ndim (.+?)\\.","errorType":"validation","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/lax.py","lineNumber":6917,"sourceCode":"\nmlir.register_lowering(\n  ragged_dot_general_p, _ragged_dot_general_gpu_lowering, platform='gpu')\n\nmlir.register_lowering(\n    ragged_dot_general_p, partial(_ragged_dot_general_lower, platform='tpu'),\n    platform='tpu')\n\n\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),","sourceCodeStart":6899,"sourceCodeEnd":6935,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/lax.py#L6899-L6935","documentation":"broadcast_in_dim requires broadcast_dimensions (a.k.a. broadcast_indices) to have one entry per operand dimension, mapping each operand dim to an output dim. If the tuple length differs from operand ndim, the mapping is ill-defined and rejected.","triggerScenarios":"Calling jax.lax.broadcast_in_dim(operand, shape, broadcast_dimensions) with len(broadcast_dimensions) != np.ndim(operand), e.g. passing 2 indices for a 3-D operand.","commonSituations":"Hand-computing broadcast indices for adding dims (e.g. expanding (B,D) to (B,1,D) or (1,B,D)) and miscounting; adapting code from lax.broadcast or numpy broadcasting where no explicit index list is used.","solutions":["Make broadcast_dimensions length exactly equal operand.ndim, listing the output index each operand dim maps to","Prefer jnp.expand_dims/jnp.reshape for simple dimension insertion","Use the recipe: new axes are output dims NOT present in broadcast_dimensions; sizes must match or operand dim must be 1"],"exampleFix":"// before\nx = jnp.zeros((4, 3))\ny = lax.broadcast_in_dim(x, (4, 5, 3), (0, 1))  # wrong length\n// after\ny = lax.broadcast_in_dim(x, (4, 5, 3), (0, 2))   # maps dim0->0, dim1->2","handlingStrategy":"validation","validationCode":"import numpy as np\nassert len(broadcast_dimensions) == np.ndim(operand)","typeGuard":"def bd_valid(operand, bd, shape) -> bool:\n    return len(bd) == np.ndim(operand)","tryCatchPattern":null,"preventionTips":["Prefer jnp.expand_dims for simple cases","Write a make_broadcast_indices(operand_shape, out_shape) helper"],"tags":["jax","broadcast-in-dim","shape-validation"],"backgroundTag":"invalid-broadcast-dimensions","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}