jax-ml/jax · error · ValueError

dimensions outside range [0, ndim): {dimensions}

Error message

dimensions outside range [0, ndim): {dimensions}

What it means

Every dimension passed to jax.lax.squeeze must satisfy 0 <= d < ndim. This ValueError reports the full list when any dimension falls outside — lax.squeeze does not support negative indices or indices >= rank (unlike jnp.squeeze's axis handling in some cases).

Source

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

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

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

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Validate/clamp: dims = [d % x.ndim for d in dimensions]
2. Recompute dims from the current shape rather than reusing stale constants
3. Prefer jnp.squeeze(x, axis=...) which accepts -1-style semantics where supported

Example fix

# before
y = jax.lax.squeeze(x, dimensions=[-1])
# after
dims = [d % x.ndim for d in [-1]]
y = jax.lax.squeeze(x, dimensions=dims)

dimensions = [d % x.ndim for d in dimensions]
assert all(0 <= d < x.ndim for d in dimensions)

def dims_in_range(dims, ndim) -> bool:
    return all(0 <= d < ndim for d in dims)

When it happens

Trigger: jax.lax.squeeze(x_2d, dimensions=[2]) or dimensions=[-1]; computed axis equal to ndim after the array lost a dim.

Common situations: Hardcoded squeeze dims after upstream reshape/squeeze reduced rank; negative-axis habit from NumPy carried into lax; loop unrolling where axis variable overshoots.

Related errors

• axis {} is out of bounds for array of shape {}
• concatenate dimension out of bounds: dimension {} for shapes
• dimensions are not unique: {dimensions}
• cannot select an axis to squeeze out which has size not equa
• Only arrays with 32-bit element types can be converted to sc

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