{"record":{"id":"0a685cdad09f0e08","repo":"jax-ml/jax","slug":"repeated-axis-in-lax-expand-dims-dimensions","errorCode":null,"errorMessage":"repeated axis in lax.expand_dims: {dimensions}","messagePattern":"repeated axis in lax\\.expand_dims: (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/lax.py","lineNumber":3835,"sourceCode":"    (3,)\n\n    Attempting to squeeze a non-unit axis results in an error:\n\n    >>> jax.lax.squeeze(x, dimensions=(0,)) # doctest: +IGNORE_EXCEPTION_DETAIL\n    Traceback (most recent call last):\n      ...\n    ValueError: cannot select an axis to squeeze out which has size not equal to one, got shape=(3, 1, 1) and dimensions=(0,)\n  \"\"\"\n  ndim = np.ndim(array)\n  dimensions = tuple(sorted(canonicalize_axis(i, ndim) for i in dimensions))\n  if not dimensions and isinstance(array, Array):\n    return array\n  return squeeze_p.bind(array, dimensions=dimensions)\n\ndef expand_dims(array: ArrayLike, dimensions: Sequence[int]) -> Array:\n  \"\"\"Insert any number of size 1 dimensions into an array.\"\"\"\n  if len(set(dimensions)) != len(dimensions):\n    raise ValueError(f'repeated axis in lax.expand_dims: {dimensions}')\n  ndim_out = np.ndim(array) + len(dimensions)\n  dims = [canonicalize_axis(i, ndim_out) for i in dimensions]\n  if len(set(dims)) != len(dims):  # check again after canonicalizing\n    raise ValueError(f'repeated axis in lax.expand_dims: {dims}')\n  dims_set = frozenset(dims)\n  result_shape = list(np.shape(array))\n  for i in sorted(dims_set):\n    result_shape.insert(i, 1)\n  broadcast_dims = [i for i in range(ndim_out) if i not in dims_set]\n  return broadcast_in_dim(array, result_shape, broadcast_dims)\n\n\n### convenience wrappers around traceables\n\ndef full_like(x: ArrayLike | DuckTypedArray,\n              fill_value: ArrayLike, dtype: DTypeLike | None = None,\n              shape: Shape | None = None, sharding: Sharding | None = None) -> Array:\n  \"\"\"Create a full array like np.full based on the example array `x`.","sourceCodeStart":3817,"sourceCodeEnd":3853,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/lax.py#L3817-L3853","documentation":"jax.lax.expand_dims rejects dimension lists containing duplicate axes. This first check fires on the raw, non-canonicalized dimensions argument, e.g. [2, 2] or [0, -1] where both resolve to the same axis only after normalization.","triggerScenarios":"Calling lax.expand_dims(x, dimensions) with a repeated entry in the raw list, e.g. (1, 1) or (0, 0, 0).","commonSituations":"Building dimension lists programmatically (list multiplication like [d]*n), copy-paste errors, or negative indices like (-1, 1) that duplicate after canonicalization (that variant raises the sibling message).","solutions":["De-duplicate the dimensions list before calling: sorted(set(dimensions))","Audit programmatic dimension construction (e.g. [d] * count bugs)","Prefer jnp.expand_dims or x[:, None, ...] style indexing for single-axis insertion"],"exampleFix":"// before\ny = lax.expand_dims(x, (1, 1))\n// after\ny = lax.expand_dims(x, (1,))\n# or y = x[:, None]","handlingStrategy":"validation","validationCode":"dims = tuple(dict.fromkeys(dimensions))  # dedupe, keep order\nndim_out = np.ndim(array) + len(dims)\ndims = tuple(sorted({canonicalize_axis(d, ndim_out) for d in dims}))\nout = lax.expand_dims(array, dims)","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Never build dimension lists by repetition ([d]*n)","Canonicalize and dedupe axes in one helper before lax.expand_dims","Prefer x[:, None] indexing for single new axes"],"tags":["jax","lax","expand-dims","duplicate-axis","value-error"],"backgroundTag":"duplicate-axis-argument","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}