{"record":{"id":"7e67dea5d9efb3cb","repo":"jax-ml/jax","slug":"dimensions-outside-range-0-ndim-dimensions","errorCode":null,"errorMessage":"dimensions outside range [0, ndim): {dimensions}","messagePattern":"dimensions outside range \\[0, ndim\\): (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/lax.py","lineNumber":7745,"sourceCode":"\ndef _squeeze_sharding_rule(operand, *, dimensions):\n  dims_set = set(dimensions)\n  new_spec = tuple(s for i, s in enumerate(operand.sharding.spec.partitions)\n                   if i not in dims_set)\n  return operand.sharding.update(\n      spec=operand.sharding.spec.update(partitions=new_spec))\n\ndef _squeeze_ur_rule(operand, *, dimensions):\n  out_unreduced = core.getu(operand)\n  kind = UnreducedKind.sum if out_unreduced else None\n  return out_unreduced, core.getr(operand), kind\n\ndef _compute_squeeze_shape(shape, dimensions):\n  dims_set = set(dimensions)\n  if len(dims_set) != len(dimensions):\n    raise ValueError(f\"dimensions are not unique: {dimensions}\")\n  if not all(0 <= d < len(shape) for d in dims_set):\n    raise ValueError(f\"dimensions outside range [0, ndim): {dimensions}\")\n  if any(not core.definitely_equal(shape[d], 1) for d in dimensions):\n    raise ValueError(\n        \"cannot select an axis to squeeze out which has size not equal to \"\n        f\"one, got {shape=} and {dimensions=}\")\n  return tuple(s for i, s in enumerate(shape) if i not in dims_set)\n\ndef _squeeze_transpose_rule(t, operand, *, dimensions):\n  assert ad.is_undefined_primal(operand)\n  return [expand_dims(t, dimensions)]\n\ndef _squeeze_batch_rule(batched_args, batch_dims, *, dimensions):\n  operand, = batched_args\n  bdim, = batch_dims\n  operand = batching.moveaxis(operand, bdim, 0)\n  dimensions = tuple(np.add(1, dimensions))\n\n  result_shape = _compute_squeeze_shape(operand.shape, dimensions)\n  bdim_out = canonicalize_axis(0, len(result_shape))","sourceCodeStart":7727,"sourceCodeEnd":7763,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/lax.py#L7727-L7763","documentation":"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).","triggerScenarios":"jax.lax.squeeze(x_2d, dimensions=[2]) or dimensions=[-1]; computed axis equal to ndim after the array lost a dim.","commonSituations":"Hardcoded squeeze dims after upstream reshape/squeeze reduced rank; negative-axis habit from NumPy carried into lax; loop unrolling where axis variable overshoots.","solutions":["Validate/clamp: dims = [d % x.ndim for d in dimensions]","Recompute dims from the current shape rather than reusing stale constants","Prefer jnp.squeeze(x, axis=...) which accepts -1-style semantics where supported"],"exampleFix":"# before\ny = jax.lax.squeeze(x, dimensions=[-1])\n# after\ndims = [d % x.ndim for d in [-1]]\ny = jax.lax.squeeze(x, dimensions=dims)","handlingStrategy":"validation","validationCode":"dimensions = [d % x.ndim for d in dimensions]\nassert all(0 <= d < x.ndim for d in dimensions)","typeGuard":"def dims_in_range(dims, ndim) -> bool:\n    return all(0 <= d < ndim for d in dims)","tryCatchPattern":null,"preventionTips":["Normalize axes modulo ndim for lax-level ops","Recompute squeeze dims after any reshape/squeeze"],"tags":["jax","squeeze","axis-out-of-bounds"],"backgroundTag":"axis-out-of-bounds","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}