jax-ml/jax · error · ValueError

repeated axis in lax.expand_dims: {dims}

Error message

repeated axis in lax.expand_dims: {dims}

What it means

jax.lax.expand_dims canonicalizes each axis (allowing negatives) against the output rank, then re-checks for duplicates. This second ValueError fires when distinct raw indices collapse to the same axis after canonicalization, e.g. (-1, 1) on a 1-d input both becoming axis 1.

Source

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

    >>> 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`.

  Args:
    x: example array-like, used for shape and dtype information.
    fill_value: a scalar value to fill the entries of the output array.

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Rewrite all indices in canonical non-negative form and deduplicate
2. Validate dims with a helper that canonicalizes first (see validationCode)
3. Use jnp.atleast_nd + reshape for explicit control of the output shape

Example fix

// before
y = lax.expand_dims(x, (-1, 1))  # both -> axis 1
// after
y = lax.expand_dims(x, (1,))

from jax._src.core import canonicalize_axis
nd = np.ndim(array) + len(dimensions)
canon = sorted({canonicalize_axis(d, nd) for d in dimensions})
out = lax.expand_dims(array, canon)

When it happens

Trigger: Calling lax.expand_dims with mixed positive/negative indices that normalize to the same output axis, e.g. expand_dims(x, (0, -2)) where both map to axis 0 of the result.

Common situations: Mixing index conventions in generated code, refactoring code that switched from positive to negative indexing without removing the old entries, or off-by-one negative indices like -1 colliding with ndim.

Related errors

• repeated axis in lax.expand_dims: {dimensions}
• lax.platform_dependent: the '{pname}' branch must be a calla
• Use 'cuda', 'rocm', or 'oneapi' for lax.platform_dependent.
• lax.platform_dependent: the 'default' branch must be a calla
• XLA operations do not allow negative axes

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