jax-ml/jax · error · ValueError

repeated axis in lax.expand_dims: {dimensions}

Error message

repeated axis in lax.expand_dims: {dimensions}

What it means

jax.lax.expand_dims rejects dimension lists containing duplicate axes. This first check fires on the raw, non-canonicalized dimensions argument, e.g. [2, 2] or [0, -1] where both resolve to the same axis only after normalization.

Source

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

    (3,)

    Attempting to squeeze a non-unit axis results in an error:

    >>> jax.lax.squeeze(x, dimensions=(0,)) # doctest: +IGNORE_EXCEPTION_DETAIL
    Traceback (most recent call last):
      ...
    ValueError: cannot select an axis to squeeze out which has size not equal to one, got shape=(3, 1, 1) and dimensions=(0,)
  """
  ndim = np.ndim(array)
  dimensions = tuple(sorted(canonicalize_axis(i, ndim) for i in dimensions))
  if not dimensions and isinstance(array, Array):
    return array
  return squeeze_p.bind(array, dimensions=dimensions)

def expand_dims(array: ArrayLike, dimensions: Sequence[int]) -> Array:
  """Insert any number of size 1 dimensions into an array."""
  if len(set(dimensions)) != len(dimensions):
    raise ValueError(f'repeated axis in lax.expand_dims: {dimensions}')
  ndim_out = np.ndim(array) + len(dimensions)
  dims = [canonicalize_axis(i, ndim_out) for i in dimensions]
  if len(set(dims)) != len(dims):  # check again after canonicalizing
    raise ValueError(f'repeated axis in lax.expand_dims: {dims}')
  dims_set = frozenset(dims)
  result_shape = list(np.shape(array))
  for i in sorted(dims_set):
    result_shape.insert(i, 1)
  broadcast_dims = [i for i in range(ndim_out) if i not in dims_set]
  return broadcast_in_dim(array, result_shape, broadcast_dims)


### convenience wrappers around traceables

def full_like(x: ArrayLike | DuckTypedArray,
              fill_value: ArrayLike, dtype: DTypeLike | None = None,
              shape: Shape | None = None, sharding: Sharding | None = None) -> Array:
  """Create a full array like np.full based on the example array `x`.

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. De-duplicate the dimensions list before calling: sorted(set(dimensions))
2. Audit programmatic dimension construction (e.g. [d] * count bugs)
3. Prefer jnp.expand_dims or x[:, None, ...] style indexing for single-axis insertion

Example fix

// before
y = lax.expand_dims(x, (1, 1))
// after
y = lax.expand_dims(x, (1,))
# or y = x[:, None]

dims = tuple(dict.fromkeys(dimensions))  # dedupe, keep order
ndim_out = np.ndim(array) + len(dims)
dims = tuple(sorted({canonicalize_axis(d, ndim_out) for d in dims}))
out = lax.expand_dims(array, dims)

When it happens

Trigger: Calling lax.expand_dims(x, dimensions) with a repeated entry in the raw list, e.g. (1, 1) or (0, 0, 0).

Common situations: Building dimension lists programmatically (list multiplication like [d]*n), copy-paste errors, or negative indices like (-1, 1) that duplicate after canonicalization (that variant raises the sibling message).

Related errors

• repeated axis in lax.expand_dims: {dims}
• Invalid dimension range passed to collapse: {operand.shape}[
• Arguments to batch_matmul must be at least 2D, got {}, {}
• Arguments to batch_matmul must have same ndim, got {}, {}
• lax.platform_dependent: the '{pname}' branch must be a calla

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