jax-ml/jax · error · ValueError

dimensions are not unique: {dimensions}

Error message

dimensions are not unique: {dimensions}

What it means

jax.lax.squeeze takes a list of dimensions to remove; listing the same dimension twice (e.g. [1, 1]) is ambiguous and raises this ValueError in _compute_squeeze_shape before any shape math is done. Duplicates could come from user lists or dimension-derivation code.

Source

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

def _squeeze_shape_rule(operand, *, dimensions):
  return _compute_squeeze_shape(np.shape(operand), dimensions)

def _squeeze_sharding_rule(operand, *, dimensions):
  dims_set = set(dimensions)
  new_spec = tuple(s for i, s in enumerate(operand.sharding.spec.partitions)
                   if i not in dims_set)
  return operand.sharding.update(
      spec=operand.sharding.spec.update(partitions=new_spec))

def _squeeze_ur_rule(operand, *, dimensions):
  out_unreduced = core.getu(operand)
  kind = UnreducedKind.sum if out_unreduced else None
  return out_unreduced, core.getr(operand), kind

def _compute_squeeze_shape(shape, dimensions):
  dims_set = set(dimensions)
  if len(dims_set) != len(dimensions):
    raise ValueError(f"dimensions are not unique: {dimensions}")
  if not all(0 <= d < len(shape) for d in dims_set):
    raise ValueError(f"dimensions outside range [0, ndim): {dimensions}")
  if any(not core.definitely_equal(shape[d], 1) for d in dimensions):
    raise ValueError(
        "cannot select an axis to squeeze out which has size not equal to "
        f"one, got {shape=} and {dimensions=}")
  return tuple(s for i, s in enumerate(shape) if i not in dims_set)

def _squeeze_transpose_rule(t, operand, *, dimensions):
  assert ad.is_undefined_primal(operand)
  return [expand_dims(t, dimensions)]

def _squeeze_batch_rule(batched_args, batch_dims, *, dimensions):
  operand, = batched_args
  bdim, = batch_dims
  operand = batching.moveaxis(operand, bdim, 0)
  dimensions = tuple(np.add(1, dimensions))

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Dedupe before calling: dimensions = list(set(dimensions)) (order does not matter for squeeze)
2. Use jnp.squeeze(x, axis=k) for a single axis
3. Compute dimension lists once from shape rather than merging ad-hoc lists

Example fix

# before
y = jax.lax.squeeze(x, dimensions=axes_a + axes_b)  # may overlap
# after
y = jax.lax.squeeze(x, dimensions=list(set(axes_a + axes_b)))

dimensions = list(set(dimensions))

def unique_dims(dims) -> bool:
    return len(set(dims)) == len(dims)

When it happens

Trigger: jax.lax.squeeze(x, dimensions=[1, 1]); concatenating dimension lists that overlap; dimensions computed from a mask/argwhere with repeats.

Common situations: Building the dims list from multiple sources (e.g. axis + extra_axes) without dedup; passing a Python range plus a manual axis that overlaps; refactors merging squeeze calls.

Related errors

• {full_name} must be a pytree prefix with bool leaves or a tu
• Axes mentioned in `manual_axis_type` field of ShapedArray sh
• varying and unreduced cannot have common mesh axes. Got vary
• varying and reduced cannot have common mesh axes. Got varyin
• `unreduced_kind` should be `None` when `unreduced` is an emp

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