jax-ml/jax · error · TypeError

broadcast_in_dim broadcast_dimensions must have length equal

Error message

broadcast_in_dim broadcast_dimensions must have length equal to operand ndim; got broadcast_dimensions {} for operand ndim {}.

What it means

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.

Source

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

mlir.register_lowering(
  ragged_dot_general_p, _ragged_dot_general_gpu_lowering, platform='gpu')

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),

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Make broadcast_dimensions length exactly equal operand.ndim, listing the output index each operand dim maps to
2. Prefer jnp.expand_dims/jnp.reshape for simple dimension insertion
3. Use the recipe: new axes are output dims NOT present in broadcast_dimensions; sizes must match or operand dim must be 1

Example fix

// before
x = jnp.zeros((4, 3))
y = lax.broadcast_in_dim(x, (4, 5, 3), (0, 1))  # wrong length
// after
y = lax.broadcast_in_dim(x, (4, 5, 3), (0, 2))   # maps dim0->0, dim1->2

import numpy as np
assert len(broadcast_dimensions) == np.ndim(operand)

def bd_valid(operand, bd, shape) -> bool:
    return len(bd) == np.ndim(operand)

When it happens

Trigger: 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.

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

Related errors

• scan got `length` argument of {} which disagrees with leadin
• conv_general_dilated batch_group_count must divide lhs batch
• conv_general_dilated rhs output feature dimension size must
• conv_general_dilated window and window_strides must have the
• Wrong number of explicit pads for convolution: expected {},

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