jax-ml/jax · error · ValueError

mode must be 'largest' or 'smallest', got {mode!r}

Error message

mode must be 'largest' or 'smallest', got {mode!r}

What it means

jnp.top_k requires mode to be exactly the string 'largest' or 'smallest'; any other value raises ValueError with the offending value echoed. The validation happens after the complex-dtype check and before axis canonicalization, so it fires even for valid arrays.

Source

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

    >>> indices
    Array([[4, 3],
           [0, 1]], dtype=int32)

    Find the two smallest elements along the first axis:

    >>> values, indices = jnp.top_k(a, 2, axis=0, mode='smallest')
    >>> values
    Array([[1, 2, 3, 2, 1],
           [5, 4, 3, 4, 5]], dtype=int32)
    >>> indices
    Array([[0, 0, 0, 1, 1],
           [1, 1, 1, 0, 0]], dtype=int32)
  """
  arr = util.ensure_arraylike("top_k", a)
  if dtypes.issubdtype(arr.dtype, np.complexfloating):
    raise ValueError("top_k is not compatible with complex inputs.")
  if mode not in ("largest", "smallest"):
    raise ValueError(f"mode must be 'largest' or 'smallest', got {mode!r}")
  axis = canonicalize_axis(axis, arr.ndim)
  if mode == "largest":
    return lax.top_k(arr, k, axis=axis)
  elif dtypes.isdtype(arr.dtype, "bool"):
    inv = lax.bitwise_not(arr)
    vals, indices = lax.top_k(inv, k, axis=axis)
    return lax.bitwise_not(vals), indices
  elif dtypes.isdtype(arr.dtype, "unsigned integer"):
    inv = -(arr + 1)
    vals, indices = lax.top_k(inv, k, axis=axis)
    return -(vals + 1), indices
  else:
    inv = -arr
    vals, indices = lax.top_k(inv, k, axis=axis)
    return -vals, indices

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use exactly 'largest' or 'smallest' as the mode string
2. Map external mode names at the boundary: {'max': 'largest', 'min': 'smallest'}
3. Validate/normalize mode strings from config before passing them in

Example fix

# before
vals, idx = jnp.top_k(x, 5, mode='max')
# after
vals, idx = jnp.top_k(x, 5, mode='largest')

assert mode in ('largest', 'smallest'), f"invalid top_k mode: {mode!r}"

def is_valid_top_k_mode(mode: str) -> bool:
    return mode in ('largest', 'smallest')

When it happens

Trigger: Calling jnp.top_k(a, k, mode='max'), mode='top', mode=None, or a typo like 'smalest'; also programmatic mode strings like 'top'/'bottom' that don't match the API.

Common situations: Assuming the API mirrors a different library's mode names (e.g., 'max'/'min' or 'top'/'bottom'); building mode from user config or CLI flags without validating against the allowed set; silent typos.

Related errors

• top_k is not compatible with complex inputs.
• Unknown resize method "{s}"
• Use 'cuda', 'rocm', or 'oneapi' for lax.platform_dependent.
• repeated axis in lax.expand_dims: {dimensions}
• Invalid dimension range passed to collapse: {operand.shape}[

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