jax-ml/jax · error · TypeError

broadcast_in_dim operand dimension sizes must either be 1, o

Error message

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 {} 

What it means

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.

Source

Thrown at jax/_src/lax/lax.py:6934

           'operand ndim; got broadcast_dimensions {} for operand ndim {}.')
    raise TypeError(msg.format(broadcast_dimensions, operand_ndim))
  if len(shape) < operand_ndim:
    msg = ('broadcast_in_dim target broadcast shape must have equal or higher rank '
           'to the operand shape; got operand ndim {} and target broadcast ndim {}.')
    raise TypeError(msg.format(operand_ndim, len(shape)))
  if not set(broadcast_dimensions).issubset(set(range(len(shape)))):
    msg = ('broadcast_in_dim broadcast_dimensions must be a subset of output '
           'dimensions, got {} for operand ndim {} and shape {}.')
    raise TypeError(msg.format(broadcast_dimensions, operand_ndim, shape))
  if not all(core.definitely_equal_one_of_dim(operand.shape[i],
                                              [1, shape[broadcast_dimensions[i]]])
             for i in range(operand_ndim)):
    msg = (
        "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 {} ")
    raise TypeError(msg.format(
        tuple(core.replace_tracer_for_error_message(d) for d in operand.shape),
        shape, broadcast_dimensions))
  if len(broadcast_dimensions) != len(set(broadcast_dimensions)):
    msg = ("broadcast_in_dim broadcast_dimensions must not contain duplicates, "
           "got broadcast_dimensions {}")
    raise TypeError(msg.format(broadcast_dimensions))
  return shape

def _broadcast_in_dim_sharding_rule(operand, *, shape, broadcast_dimensions,
                                    sharding):
  if sharding is not None:
    return sharding
  bds = set(broadcast_dimensions)
  orig_spec = iter(operand.sharding.spec.partitions)
  new_spec = [next(orig_spec) if i in bds else None for i in range(len(shape))]
  assert next(orig_spec, None) is None
  mesh = (get_abstract_mesh() if operand.sharding.mesh.empty else
          operand.sharding.mesh)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Fix the target shape so dim sizes match the operand, or make the operand dim 1 before broadcasting
2. Use lax.tile / jnp.tile for repetition of non-unit dims
3. Check the mapping index: often the wrong output index pairs the operand dim with an unrelated dim

Example fix

// before
x = jnp.zeros((4, 3))
y = lax.broadcast_in_dim(x, (4, 5), (0, 1))  # operand dim 3 vs target 5
// after
x1 = x[:, None]                              # (4, 1)
y = lax.broadcast_in_dim(x1, (4, 5), (0, 1))  # 1 broadcasts to 5

ok = all(operand.shape[i] in (1, shape[bd[i]]) for i in range(np.ndim(operand)))
assert ok

def can_broadcast(operand_shape, shape, bd) -> bool:
    return all(operand_shape[i] == 1 or operand_shape[i] == shape[bd[i]]
               for i in range(len(operand_shape)))

When it happens

Trigger: Calling broadcast_in_dim where operand.shape[i] != shape[broadcast_dimensions[i]] and operand.shape[i] != 1.

Common situations: 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.

Related errors

• unexpected JAX type (e.g. shape/dtype) for argument to VJP f
• cotangent type does not match function output, expected {out
• Mismatched number of outputs from callback. Expected: {}, Ac
• Incorrect output shape for return value #{i}: Expected: {out
• Length of sharding.spec ({len(out_s.spec)}) must be equal to

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/4640ced5fd30d6bd. Report an issue: GitHub.