jax-ml/jax · error · TypeError

Arguments to sort must have equal shapes, got

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.

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…
• dynamic_update_slice update shape must be smaller than…
• A must be ( , ) matrix A, got output
• Accumulator and LHS have incompatible shapes. Expected LHS…
• Accumulator and RHS have incompatible shapes. Expected RHS…

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

Appendix: 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)