{"record":{"id":"d870c0da9b5f95a5","repo":"jax-ml/jax","slug":"order-argument-to-argsort-is-not-supported","errorCode":null,"errorMessage":"'order' argument to argsort is not supported.","messagePattern":"'order' argument to argsort is not supported\\.","errorType":"validation","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/numpy/sorting.py","lineNumber":155,"sourceCode":"    >>> indices\n    Array([[1, 0, 2],\n           [2, 1, 0]], dtype=int32)\n    >>> jnp.take_along_axis(x, indices, axis=1)\n    Array([[1, 2, 3],\n           [3, 4, 6]], dtype=int32)\n\n\n  See also:\n    - :func:`jax.numpy.sort`: return sorted values directly.\n    - :func:`jax.numpy.lexsort`: lexicographical sort of multiple arrays.\n    - :func:`jax.lax.sort`: lower-level function wrapping XLA's Sort operator.\n  \"\"\"\n  arr = util.ensure_arraylike(\"argsort\", a)\n  if kind is not None:\n    raise TypeError(\"'kind' argument to argsort is not supported. Use\"\n                    \" stable=True or stable=False to specify sort stability.\")\n  if order is not None:\n    raise TypeError(\"'order' argument to argsort is not supported.\")\n  if axis is None:\n    arr = arr.ravel()\n    axis = 0\n  dimension = canonicalize_axis(axis, arr.ndim)\n  if dtype is not None:\n    idx_dtype = dtypes.check_and_canonicalize_user_dtype(dtype, \"argsort\")\n  else:\n    idx_dtype = lax_utils.int_dtype_for_dim(arr.shape[dimension], 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, arr.shape, dimension,\n                              out_sharding=core.typeof(arr).sharding)\n  # For stable descending sort, we reverse the array and indices to ensure that\n  # duplicates remain in their original order when the final indices are reversed.\n  # For non-stable descending sort, we can avoid these extra operations.\n  if descending and stable:","sourceCodeStart":137,"sourceCodeEnd":173,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/numpy/sorting.py#L137-L173","documentation":"jax.numpy.argsort explicitly rejects the NumPy 'order' argument because JAX arrays do not support structured/record dtypes that 'order' sorts by. JAX reimplements NumPy's API surface but only supports the subset meaningful for XLA compilation; 'kind' is likewise rejected in favor of the 'stable' boolean. Passing 'order' (even order=None explicitly is fine, but any non-None value) raises TypeError.","triggerScenarios":"Calling jnp.argsort(a, order=['x']) or with any non-None order keyword, e.g. jnp.argsort(a, order=0). Also triggered when porting NumPy code that sorts structured arrays by field names.","commonSituations":"Porting NumPy/pandas sorting code to JAX; using structured arrays (recarray) which JAX doesn't support; copy-pasting np.argsort calls with field ordering from a data-processing pipeline.","solutions":["Remove the order= argument from the jnp.argsort call","If sorting by a field of a structured array, decompose the data into separate arrays (e.g., pytrees or dict of arrays) and argsort the key array directly","Replace kind='stable' with stable=True (related restriction in the same function)","If you truly need structured sorting, do it in NumPy on the host before transferring to device"],"exampleFix":"// before\nidx = jnp.argsort(data, order=('score', 'id'))\n// after\nidx = jnp.argsort(scores)  # sort by the extracted key array","handlingStrategy":"validation","validationCode":"assert order is None, \"jnp.argsort does not support order; sort key arrays directly\"","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Strip NumPy-only kwargs (order, kind) when porting to jnp","Use stable=True/False instead of kind","Keep structured-array logic on the NumPy side"],"tags":["jax","argsort","numpy-compat","unsupported-argument","sorting"],"backgroundTag":"unsupported-argument-compatibility","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}