jax-ml/jax · error · TypeError

need sequence of keys with len > 0 in lexsort

Error message

need sequence of keys with len > 0 in lexsort

What it means

jnp.lexsort requires at least one sort key; passing an empty sequence is a TypeError because there is nothing to define an ordering. NumPy similarly errors, and JAX mirrors that contract. The check fires on len(keys) == 0 after input conversion.

Source

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

    >>> key1 = jnp.array([[2, 4, 2, 3],
    ...                   [3, 1, 2, 2]])
    >>> key2 = jnp.array([[1, 2, 1, 3],
    ...                   [2, 1, 2, 1]])
    >>> jnp.lexsort([key1, key2])
    Array([[0, 2, 1, 3],
           [1, 3, 2, 0]], dtype=int32)

    A different sort axis can be chosen using the ``axis`` keyword; here we sort
    along the leading axis:

    >>> jnp.lexsort([key1, key2], axis=0)
    Array([[0, 1, 0, 1],
           [1, 0, 1, 0]], dtype=int32)
  """
  key_arrays = util.ensure_arraylike_tuple("lexsort", tuple(keys))
  if len(key_arrays) == 0:
    raise TypeError("need sequence of keys with len > 0 in lexsort")
  if len({np.shape(key) for key in key_arrays}) > 1:
    raise ValueError("all keys need to be the same shape")
  if np.ndim(key_arrays[0]) == 0:
    return lax.full((), 0, dtypes.default_int_dtype())
  axis = canonicalize_axis(axis, np.ndim(key_arrays[0]))
  idx_dtype = lax_utils.int_dtype_for_dim(key_arrays[0].shape[axis],
                                          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, np.shape(key_arrays[0]), axis)
  return lax.sort((*key_arrays[::-1], iota), dimension=axis, num_keys=len(key_arrays))[-1]


@export
@api.jit(static_argnums=1, static_argnames=('axis', 'mode', 'sorted'))
def top_k(

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Guard the call: only invoke lexsort when the keys list is non-empty
2. Provide a fallback ordering when no keys exist (e.g., use arange for identity order)
3. Fix the upstream logic that produced an empty key list

Example fix

# before
order = jnp.lexsort(keys)  # keys may be []
# after
order = jnp.lexsort(keys) if keys else jnp.arange(n)

assert len(keys) > 0, 'lexsort needs at least one key'

def valid_lexsort_keys(keys) -> bool:
    return len(keys) > 0

When it happens

Trigger: Calling jnp.lexsort([]) or jnp.lexsort(tuple()) — typically because a keys list was built dynamically (e.g., [k for k in ... if pred]) and ended up empty.

Common situations: Programmatically building a key list from user input or config that can be empty; refactoring multi-key sorts so the key list is computed rather than hardcoded; default-argument bugs where keys=None is converted to an empty tuple.

Related errors

• all keys need to be the same shape
• Need at least one array to concatenate
• {full_name} must be a pytree prefix with bool leaves or a tu
• Axes mentioned in `manual_axis_type` field of ShapedArray sh
• varying and unreduced cannot have common mesh axes. Got vary

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