jax-ml/jax · error · TypeError

broadcast_in_dim broadcast_dimensions must not contain dupli

Error message

broadcast_in_dim broadcast_dimensions must not contain duplicates, got broadcast_dimensions {}

What it means

broadcast_dimensions must not repeat an index: each operand dim must map to a distinct output dim. Duplicate indices would make the mapping non-injective and the broadcast ambiguous.

Source

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

  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)
  return operand.sharding.update(
      mesh=mesh, spec=operand.sharding.spec.update(partitions=new_spec))

def _broadcast_in_dim_unreduced_rule(operand, sharding):
  if sharding is not None and sharding.mesh.are_all_axes_explicit:
    out = sharding.spec.unreduced

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Remove duplicates so each operand dim maps to a unique output dim
2. Build indices programmatically and assert len(set(bd)) == len(bd) before calling
3. Reconsider which output dims you actually intend each operand dim to align with

Example fix

// before
x = jnp.zeros((2, 3))
y = lax.broadcast_in_dim(x, (2, 3, 4), (0, 0))
// after
y = lax.broadcast_in_dim(x, (2, 3, 4), (0, 1))

assert len(set(broadcast_dimensions)) == len(broadcast_dimensions)

When it happens

Trigger: Calling broadcast_in_dim with e.g. broadcast_dimensions=(0, 0) or (1, 1, 2).

Common situations: Typing the same index twice when constructing index tuples by hand; generating indices with a buggy comprehension.

Related errors

• broadcast_in_dim broadcast_dimensions must have length equal
• broadcast_in_dim target broadcast shape must have equal or h
• broadcast_in_dim broadcast_dimensions must be a subset of ou
• broadcast_in_dim operand dimension sizes must either be 1, o
• `out_sharding` passed to `broadcast_in_dim` can only contain

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