{"record":{"id":"b735649eec11c1c9","repo":"jax-ml/jax","slug":"cannot-select-an-axis-to-squeeze-out-which-has-siz","errorCode":null,"errorMessage":"cannot select an axis to squeeze out which has size not equal to one, got {shape=} and {dimensions=}","messagePattern":"cannot select an axis to squeeze out which has size not equal to one, got (.+?) and (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/lax.py","lineNumber":7747,"sourceCode":"  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))\n  return squeeze(operand, dimensions=dimensions), bdim_out\n","sourceCodeStart":7729,"sourceCodeEnd":7765,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/lax.py#L7729-L7765","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Conditionally squeeze only size-1 axes: dims = [i for i, s in enumerate(x.shape) if s == 1 and i in wanted]","Use reshape to a target shape instead of squeeze when sizes are known: x.reshape(...)","Keep reductions with keepdims=False instead of manually squeezing"],"exampleFix":"# before\ny = jax.lax.squeeze(x, dimensions=[0])  # fails when batch > 1\n# after\nif x.shape[0] == 1:\n    y = jax.lax.squeeze(x, dimensions=[0])\nelse:\n    y = x","handlingStrategy":"type-guard","validationCode":"dims = [d for d in dimensions\n        if d < x.ndim and getattr(x.shape[d], 'value', x.shape[d]) == 1]\nout = jax.lax.squeeze(x, dimensions=dims) if dims else x","typeGuard":"def squeezable(x, dims) -> bool:\n    return all(0 <= d < x.ndim and x.shape[d] == 1 for d in dims)","tryCatchPattern":null,"preventionTips":["Filter to size-1 axes before squeezing, especially under jit with dynamic dims","Use reshape with explicit target shapes when sizes are known"],"tags":["jax","squeeze","size-one-required","shape-validation"],"backgroundTag":"squeeze-non-unit-axis","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}