jax-ml/jax · error · ValueError

all keys need to be the same shape

Error message

all keys need to be the same shape

What it means

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.

Source

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

    ...                   [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(
    a: ArrayLike,
    k: int,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Check and align key shapes before the call: assert all(k.shape == keys[0].shape for k in keys)
2. Reshape or broadcast keys to a common shape deliberately (e.g., jnp.broadcast_arrays) if that is semantically correct
3. Fix the upstream code that altered one key's shape (remove stray ravel/reshape)
4. Verify you are not accidentally passing rows vs columns (transpose) for one key

Example fix

# before
order = jnp.lexsort([secondary, primary])  # shapes (3,) and (3, 1)
# after
secondary, primary = jnp.broadcast_arrays(secondary, primary)
order = jnp.lexsort([secondary, primary])

shapes = {k.shape for k in keys}
assert len(shapes) == 1, f'lexsort key shapes differ: {shapes}'

def lexsort_keys_aligned(keys) -> bool:
    s = np.shape(keys[0])
    return all(np.shape(k) == s for k in keys)

When it happens

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

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

Related errors

• need sequence of keys with len > 0 in lexsort
• Sum of sizes {n} must be equal to dimension {axis} of the op
• grid_names must have the same length as grid, got {self}.
• cluster_names must have the same length as cluster, got {sel
• All arrays must have matching shapes along non-concatenated

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