jax-ml/jax · error · TypeError

broadcast_in_dim target broadcast shape must have equal or h

Error message

broadcast_in_dim target broadcast shape must have equal or higher rank to the operand shape; got operand ndim {} and target broadcast ndim {}.

What it means

broadcast_in_dim can only add leading/trailing-style new dimensions, so the target shape must have rank >= operand rank. A lower-rank target is impossible to map and raises this TypeError.

Source

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

mlir.register_lowering(
    ragged_dot_general_p, partial(_ragged_dot_general_lower, platform='tpu'),
    platform='tpu')


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

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Ensure the target shape rank is >= operand rank; broadcasting never removes dimensions
2. Use lax.reshape or slicing to reduce rank before broadcasting
3. Double-check the intended output shape arithmetic in your code

Example fix

// before
x = jnp.zeros((2, 3))
y = lax.broadcast_in_dim(x, (2,), (0, 1))  # target rank 1 < operand rank 2
// after
y = lax.broadcast_in_dim(x, (2, 3, 4), (0, 1))  # rank 3 >= 2

assert len(shape) >= np.ndim(operand), 'broadcast can only add dims'

When it happens

Trigger: Calling broadcast_in_dim(operand, shape, ...) where len(shape) < np.ndim(operand).

Common situations: Passing an already-higher-rank operand with a scalar/low-rank shape (e.g. targeting shape=() with a vector); confusion between broadcasting up vs. reshaping down.

Related errors

• process_id and num_processes must be nonnegative, with proce
• broadcast_in_dim broadcast_dimensions must have length equal
• broadcast_in_dim broadcast_dimensions must be a subset of ou
• broadcast_in_dim operand dimension sizes must either be 1, o
• broadcast_in_dim broadcast_dimensions must not contain dupli

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