{"record":{"id":"f3d2ff8eff3b1ae8","repo":"jax-ml/jax","slug":"arguments-to-sort-must-have-equal-shapes-got-sh","errorCode":null,"errorMessage":"Arguments to sort must have equal shapes, got: {shapes}","messagePattern":"Arguments to sort must have equal shapes, got: (.+?)","errorType":"validation","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/lax.py","lineNumber":8867,"sourceCode":"\n_UINT_DTYPES = {\n  16: np.dtype(np.uint16),\n  32: np.dtype(np.uint32),\n  64: np.dtype(np.uint64),\n}\n\n_INT_DTYPES = {\n  16: np.dtype(np.int16),\n  32: np.dtype(np.int32),\n  64: np.dtype(np.int64),\n}\n\n\ndef _sort_abstract_eval(*avals, dimension, is_stable, num_keys):\n  avals = tuple(avals)\n  if any(arg.shape != avals[0].shape for arg in avals[1:]):\n    shapes = \" \".join(str(a.shape) for a in avals)\n    raise TypeError(f\"Arguments to sort must have equal shapes, got: {shapes}\")\n  non_empty_s = [\n      a.sharding for a in avals\n      if not a.sharding.mesh.empty and a.sharding.mesh._any_axis_explicit]\n  for s in non_empty_s:\n    if s.spec[dimension] is not None:\n      raise core.ShardingTypeError(\n          \"Arguments to sort must be unsharded over the sorting dimension. \"\n          f\"Got arg sharding={s} and sorting dimension={dimension}\")\n    if s != non_empty_s[0]:\n      shardings = \" \".join(str(s) for s in non_empty_s)\n      raise core.ShardingTypeError(\n          f'Arguments to sort must have equal shardings, got: {shardings}')\n  return avals\n\n\ndef _canonicalize_float_for_sort(x):\n  # In the sort comparator, we are going to use a comparison operator where -0\n  # would be before 0, and -NaN and NaN appear at the beginning and end of the","sourceCodeStart":8849,"sourceCodeEnd":8885,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/lax.py#L8849-L8885","documentation":"sort (and argsort-derived ops) require all operands (keys and values) to have identical shapes, since they are jointly permuted along one dimension. Any shape mismatch is a TypeError from the abstract evaluator.","triggerScenarios":"lax.sort(keys, values, dimension=0) where keys.shape != values.shape; often reached when sorting a key array together with an index/value array of different length.","commonSituations":"Implementing top-k-with-values or argsort by sorting (key, index) pairs where the index array was built with the wrong length; batch-dimension mismatches after slicing keys but not values.","solutions":["Broadcast or rebuild the value array to the key shape before sorting: values = jnp.broadcast_to(values, keys.shape).","If sorting (value, index) pairs, build indices with jnp.arange(keys.size).reshape(keys.shape).","Check for an accidental extra/missing dimension from a slice (e.g., keys[:, None] vs keys)."],"exampleFix":"# before\nidx = jnp.arange(keys.shape[0])\n_, order = lax.sort(keys, idx, dimension=-1)  # keys is 2-D, idx is 1-D\n# after\nidx = jnp.broadcast_to(jnp.arange(keys.shape[-1]), keys.shape)\n_, order = lax.sort(keys, idx, dimension=-1)","handlingStrategy":"validation","validationCode":"assert keys.shape == values.shape, (keys.shape, values.shape)\nif values.shape != keys.shape:\n    values = jnp.broadcast_to(values, keys.shape)\n_, order = lax.sort(keys, values, dimension=-1)","typeGuard":"def same_shapes(*arrays):\n    return all(a.shape == arrays[0].shape for a in arrays[1:])","tryCatchPattern":null,"preventionTips":["Build index arrays with arange reshaped to the key shape.","Broadcast value arrays to keys before joint sorting."],"tags":["jax","lax","sort","shape-mismatch"],"backgroundTag":"shape-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}