jax-ml/jax · error · ValueError

cannot select an axis to squeeze out which has size not equa

Error message

cannot select an axis to squeeze out which has size not equal to one, got {shape=} and {dimensions=}

What it means

Squeeze removes only size-1 axes; if any requested dimension has a size that is not definitely 1 (including dynamic/unknown sizes), _compute_squeeze_shape raises this ValueError showing shape and dimensions. The 'definitely_equal' check means even possibly-1 dynamic dims are rejected under tracing.

Source

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

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

  result_shape = _compute_squeeze_shape(operand.shape, dimensions)
  bdim_out = canonicalize_axis(0, len(result_shape))
  return squeeze(operand, dimensions=dimensions), bdim_out

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Conditionally squeeze only size-1 axes: dims = [i for i, s in enumerate(x.shape) if s == 1 and i in wanted]
2. Use reshape to a target shape instead of squeeze when sizes are known: x.reshape(...)
3. Keep reductions with keepdims=False instead of manually squeezing

Example fix

# before
y = jax.lax.squeeze(x, dimensions=[0])  # fails when batch > 1
# after
if x.shape[0] == 1:
    y = jax.lax.squeeze(x, dimensions=[0])
else:
    y = x

dims = [d for d in dimensions
        if d < x.ndim and getattr(x.shape[d], 'value', x.shape[d]) == 1]
out = jax.lax.squeeze(x, dimensions=dims) if dims else x

def squeezable(x, dims) -> bool:
    return all(0 <= d < x.ndim and x.shape[d] == 1 for d in dims)

When it happens

Trigger: jax.lax.squeeze(jnp.zeros((2,3)), dimensions=[1]) — axis of size 3; squeezing batch axes that are dynamic under jit/vmap; squeezing after a reshape that made the axis > 1.

Common situations: Assuming a singleton batch/time axis that becomes >1 with real data; dynamic batch sizes under jit making size non-definitely-1; model refactors changing axis sizes while squeeze args stayed fixed.

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/b735649eec11c1c9. Report an issue: GitHub.