jax-ml/jax · error · TypeError

'order' argument to argsort is not supported.

Error message

'order' argument to argsort is not supported.

What it means

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.

Source

Thrown at jax/_src/numpy/sorting.py:155

    >>> indices
    Array([[1, 0, 2],
           [2, 1, 0]], dtype=int32)
    >>> jnp.take_along_axis(x, indices, axis=1)
    Array([[1, 2, 3],
           [3, 4, 6]], dtype=int32)


  See also:
    - :func:`jax.numpy.sort`: return sorted values directly.
    - :func:`jax.numpy.lexsort`: lexicographical sort of multiple arrays.
    - :func:`jax.lax.sort`: lower-level function wrapping XLA's Sort operator.
  """
  arr = util.ensure_arraylike("argsort", a)
  if kind is not None:
    raise TypeError("'kind' argument to argsort is not supported. Use"
                    " stable=True or stable=False to specify sort stability.")
  if order is not None:
    raise TypeError("'order' argument to argsort is not supported.")
  if axis is None:
    arr = arr.ravel()
    axis = 0
  dimension = canonicalize_axis(axis, arr.ndim)
  if dtype is not None:
    idx_dtype = dtypes.check_and_canonicalize_user_dtype(dtype, "argsort")
  else:
    idx_dtype = lax_utils.int_dtype_for_dim(arr.shape[dimension], signed=True)
    # We'd give the correct output values with int32, but use the default dtype to
    # match NumPy type semantics if x64 mode is enabled for now.
    if idx_dtype == np.dtype(np.int32):
      idx_dtype = dtypes.default_int_dtype()
  iota = lax.broadcasted_iota(idx_dtype, arr.shape, dimension,
                              out_sharding=core.typeof(arr).sharding)
  # For stable descending sort, we reverse the array and indices to ensure that
  # duplicates remain in their original order when the final indices are reversed.
  # For non-stable descending sort, we can avoid these extra operations.
  if descending and stable:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Remove the order= argument from the jnp.argsort call
2. 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
3. Replace kind='stable' with stable=True (related restriction in the same function)
4. If you truly need structured sorting, do it in NumPy on the host before transferring to device

Example fix

// before
idx = jnp.argsort(data, order=('score', 'id'))
// after
idx = jnp.argsort(scores)  # sort by the extracted key array

assert order is None, "jnp.argsort does not support order; sort key arrays directly"

When it happens

Trigger: 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.

Common situations: 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.

Related errors

• The 'out' argument to jnp.argmax is not supported.
• The 'out' argument to jnp.argmin is not supported.
• The 'out' argument to jnp.nanargmax is not supported.
• The 'out' argument to jnp.nanargmin is not supported.
• The 'out' argument to jnp.compress is not supported.

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/d870c0da9b5f95a5. Report an issue: GitHub.