{"record":{"id":"da7a3dcf746a8148","repo":"jax-ml/jax","slug":"repeated-axis-in-lax-expand-dims-dims","errorCode":null,"errorMessage":"repeated axis in lax.expand_dims: {dims}","messagePattern":"repeated axis in lax\\.expand_dims: (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/lax.py","lineNumber":3839,"sourceCode":"    >>> 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`.\n\n  Args:\n    x: example array-like, used for shape and dtype information.\n    fill_value: a scalar value to fill the entries of the output array.","sourceCodeStart":3821,"sourceCodeEnd":3857,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/lax.py#L3821-L3857","documentation":"jax.lax.expand_dims canonicalizes each axis (allowing negatives) against the output rank, then re-checks for duplicates. This second ValueError fires when distinct raw indices collapse to the same axis after canonicalization, e.g. (-1, 1) on a 1-d input both becoming axis 1.","triggerScenarios":"Calling lax.expand_dims with mixed positive/negative indices that normalize to the same output axis, e.g. expand_dims(x, (0, -2)) where both map to axis 0 of the result.","commonSituations":"Mixing index conventions in generated code, refactoring code that switched from positive to negative indexing without removing the old entries, or off-by-one negative indices like -1 colliding with ndim.","solutions":["Rewrite all indices in canonical non-negative form and deduplicate","Validate dims with a helper that canonicalizes first (see validationCode)","Use jnp.atleast_nd + reshape for explicit control of the output shape"],"exampleFix":"// before\ny = lax.expand_dims(x, (-1, 1))  # both -> axis 1\n// after\ny = lax.expand_dims(x, (1,))","handlingStrategy":"validation","validationCode":"from jax._src.core import canonicalize_axis\nnd = np.ndim(array) + len(dimensions)\ncanon = sorted({canonicalize_axis(d, nd) for d in dimensions})\nout = lax.expand_dims(array, canon)","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Avoid mixing positive and negative axis indices","Run axes through canonicalize_axis + set() before calling","Add a unit test asserting no duplicate axes for generated dim lists"],"tags":["jax","lax","expand-dims","axis-canonicalization","duplicate-axis"],"backgroundTag":"duplicate-axis-argument","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}