{"record":{"id":"09002b94750963b6","repo":"jax-ml/jax","slug":"all-keys-need-to-be-the-same-shape","errorCode":null,"errorMessage":"all keys need to be the same shape","messagePattern":"all keys need to be the same shape","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/numpy/sorting.py","lineNumber":432,"sourceCode":"    ...                   [3, 1, 2, 2]])\n    >>> key2 = jnp.array([[1, 2, 1, 3],\n    ...                   [2, 1, 2, 1]])\n    >>> jnp.lexsort([key1, key2])\n    Array([[0, 2, 1, 3],\n           [1, 3, 2, 0]], dtype=int32)\n\n    A different sort axis can be chosen using the ``axis`` keyword; here we sort\n    along the leading axis:\n\n    >>> jnp.lexsort([key1, key2], axis=0)\n    Array([[0, 1, 0, 1],\n           [1, 0, 1, 0]], dtype=int32)\n  \"\"\"\n  key_arrays = util.ensure_arraylike_tuple(\"lexsort\", tuple(keys))\n  if len(key_arrays) == 0:\n    raise TypeError(\"need sequence of keys with len > 0 in lexsort\")\n  if len({np.shape(key) for key in key_arrays}) > 1:\n    raise ValueError(\"all keys need to be the same shape\")\n  if np.ndim(key_arrays[0]) == 0:\n    return lax.full((), 0, dtypes.default_int_dtype())\n  axis = canonicalize_axis(axis, np.ndim(key_arrays[0]))\n  idx_dtype = lax_utils.int_dtype_for_dim(key_arrays[0].shape[axis],\n                                          signed=True)\n  # We'd give the correct output values with int32, but use the default dtype to\n  # match NumPy type semantics if x64 mode is enabled for now.\n  if idx_dtype == np.dtype(np.int32):\n    idx_dtype = dtypes.default_int_dtype()\n  iota = lax.broadcasted_iota(idx_dtype, np.shape(key_arrays[0]), axis)\n  return lax.sort((*key_arrays[::-1], iota), dimension=axis, num_keys=len(key_arrays))[-1]\n\n\n@export\n@api.jit(static_argnums=1, static_argnames=('axis', 'mode', 'sorted'))\ndef top_k(\n    a: ArrayLike,\n    k: int,","sourceCodeStart":414,"sourceCodeEnd":450,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/numpy/sorting.py#L414-L450","documentation":"All keys passed to jnp.lexsort must have identical shapes because the function produces a single permutation valid for every key (the last key is primary). Shapes are collected into a set; if more than one distinct shape exists, ValueError is raised before any axis processing.","triggerScenarios":"Calling jnp.lexsort([a, b]) where a.shape != b.shape, e.g. jnp.lexsort([jnp.zeros((3, 4)), jnp.ones((3,))]); also broadcasting mistakes where one key was squeezed or raveled differently.","commonSituations":"Multi-column sorting of tabular data where columns have mismatched lengths after preprocessing; mixing a 1-D key with a 2-D key unintentionally; a bug upstream that reshaped one key (e.g., an extra .ravel() or reshape).","solutions":["Check and align key shapes before the call: assert all(k.shape == keys[0].shape for k in keys)","Reshape or broadcast keys to a common shape deliberately (e.g., jnp.broadcast_arrays) if that is semantically correct","Fix the upstream code that altered one key's shape (remove stray ravel/reshape)","Verify you are not accidentally passing rows vs columns (transpose) for one key"],"exampleFix":"# before\norder = jnp.lexsort([secondary, primary])  # shapes (3,) and (3, 1)\n# after\nsecondary, primary = jnp.broadcast_arrays(secondary, primary)\norder = jnp.lexsort([secondary, primary])","handlingStrategy":"validation","validationCode":"shapes = {k.shape for k in keys}\nassert len(shapes) == 1, f'lexsort key shapes differ: {shapes}'","typeGuard":"def lexsort_keys_aligned(keys) -> bool:\n    s = np.shape(keys[0])\n    return all(np.shape(k) == s for k in keys)","tryCatchPattern":null,"preventionTips":["Broadcast keys with jnp.broadcast_arrays before multi-key sorts","Assert uniform column shapes in tabular preprocessing","Log shapes of keys when debugging sort pipelines"],"tags":["jax","lexsort","shape-mismatch","validation","sorting"],"backgroundTag":"shape-mismatch-validation","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}