{"record":{"id":"ce4986eec1c16c6b","repo":"jax-ml/jax","slug":"dimensions-are-not-unique-dimensions","errorCode":null,"errorMessage":"dimensions are not unique: {dimensions}","messagePattern":"dimensions are not unique: (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/lax.py","lineNumber":7743,"sourceCode":"def _squeeze_shape_rule(operand, *, dimensions):\n  return _compute_squeeze_shape(np.shape(operand), dimensions)\n\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","sourceCodeStart":7725,"sourceCodeEnd":7761,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/lax.py#L7725-L7761","documentation":"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.","triggerScenarios":"jax.lax.squeeze(x, dimensions=[1, 1]); concatenating dimension lists that overlap; dimensions computed from a mask/argwhere with repeats.","commonSituations":"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.","solutions":["Dedupe before calling: dimensions = list(set(dimensions)) (order does not matter for squeeze)","Use jnp.squeeze(x, axis=k) for a single axis","Compute dimension lists once from shape rather than merging ad-hoc lists"],"exampleFix":"# before\ny = jax.lax.squeeze(x, dimensions=axes_a + axes_b)  # may overlap\n# after\ny = jax.lax.squeeze(x, dimensions=list(set(axes_a + axes_b)))","handlingStrategy":"validation","validationCode":"dimensions = list(set(dimensions))","typeGuard":"def unique_dims(dims) -> bool:\n    return len(set(dims)) == len(dims)","tryCatchPattern":null,"preventionTips":["Dedupe merged dimension lists before squeeze","Build dims from shape introspection, not list concatenation"],"tags":["jax","squeeze","duplicate-dims","validation"],"backgroundTag":"duplicate-axis","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}