jax-ml/jax · error · TypeError

Arguments to sort must have equal shapes, got: {shapes}

Error message

Arguments to sort must have equal shapes, got: {shapes}

What it means

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.

Source

Thrown at jax/_src/lax/lax.py:8867

_UINT_DTYPES = {
  16: np.dtype(np.uint16),
  32: np.dtype(np.uint32),
  64: np.dtype(np.uint64),
}

_INT_DTYPES = {
  16: np.dtype(np.int16),
  32: np.dtype(np.int32),
  64: np.dtype(np.int64),
}


def _sort_abstract_eval(*avals, dimension, is_stable, num_keys):
  avals = tuple(avals)
  if any(arg.shape != avals[0].shape for arg in avals[1:]):
    shapes = " ".join(str(a.shape) for a in avals)
    raise TypeError(f"Arguments to sort must have equal shapes, got: {shapes}")
  non_empty_s = [
      a.sharding for a in avals
      if not a.sharding.mesh.empty and a.sharding.mesh._any_axis_explicit]
  for s in non_empty_s:
    if s.spec[dimension] is not None:
      raise core.ShardingTypeError(
          "Arguments to sort must be unsharded over the sorting dimension. "
          f"Got arg sharding={s} and sorting dimension={dimension}")
    if s != non_empty_s[0]:
      shardings = " ".join(str(s) for s in non_empty_s)
      raise core.ShardingTypeError(
          f'Arguments to sort must have equal shardings, got: {shardings}')
  return avals


def _canonicalize_float_for_sort(x):
  # In the sort comparator, we are going to use a comparison operator where -0
  # would be before 0, and -NaN and NaN appear at the beginning and end of the

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Broadcast or rebuild the value array to the key shape before sorting: values = jnp.broadcast_to(values, keys.shape).
2. If sorting (value, index) pairs, build indices with jnp.arange(keys.size).reshape(keys.shape).
3. Check for an accidental extra/missing dimension from a slice (e.g., keys[:, None] vs keys).

Example fix

# before
idx = jnp.arange(keys.shape[0])
_, order = lax.sort(keys, idx, dimension=-1)  # keys is 2-D, idx is 1-D
# after
idx = jnp.broadcast_to(jnp.arange(keys.shape[-1]), keys.shape)
_, order = lax.sort(keys, idx, dimension=-1)

assert keys.shape == values.shape, (keys.shape, values.shape)
if values.shape != keys.shape:
    values = jnp.broadcast_to(values, keys.shape)
_, order = lax.sort(keys, values, dimension=-1)

def same_shapes(*arrays):
    return all(a.shape == arrays[0].shape for a in arrays[1:])

When it happens

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

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

Related errors

• conv_general_dilated lhs and rhs must have the same number o
• dynamic_update_slice update shape must be smaller than opera
• unexpected JAX type (e.g. shape/dtype) for argument to VJP f
• cotangent type does not match function output, expected {out
• Mismatched number of outputs from callback. Expected: {}, Ac

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