jax-ml/jax · error · TypeError

broadcast_in_dim broadcast_dimensions must be a subset of ou

Error message

broadcast_in_dim broadcast_dimensions must be a subset of output dimensions, got {} for operand ndim {} and shape {}.

What it means

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.

Source

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

def _broadcast_in_dim_shape_rule(operand, *, shape, broadcast_dimensions,
                                 sharding):
  _check_shapelike('broadcast_in_dim', 'shape', shape)
  _check_shapelike('broadcast_in_dim', 'broadcast_dimensions',
                   broadcast_dimensions)
  operand_ndim = np.ndim(operand)
  if operand_ndim != len(broadcast_dimensions):
    msg = ('broadcast_in_dim broadcast_dimensions must have length equal to '
           '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,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Validate indices: all(0 <= i < len(shape) for i in broadcast_dimensions)
2. Remember indices are into the OUTPUT shape, and the max valid index is len(shape)-1
3. Write a small helper that derives broadcast_dimensions from operand/output shapes

Example fix

// before
x = jnp.zeros((3,))
y = lax.broadcast_in_dim(x, (2, 3, 4), (1, 3))  # 3 out of range for rank-3
// after
y = lax.broadcast_in_dim(x, (2, 3, 4), (1, 2))

assert all(0 <= i < len(shape) for i in broadcast_dimensions)

When it happens

Trigger: Calling broadcast_in_dim(operand, shape, broadcast_dimensions) where some index >= len(shape) or < 0 (e.g. index 3 for a 3-D output).

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

Related errors

• dot_general requires lhs dimension numbers to be nonnegative
• dot_general requires rhs dimension numbers to be nonnegative
• ragged_dot_general requires {dim_name} numbers to be nonnega
• broadcast_in_dim broadcast_dimensions must have length equal
• broadcast_in_dim target broadcast shape must have equal or h

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