jax-ml/jax · error · ValueError

mode must be 'largest' or 'smallest', got

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.

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.

Understand the failure class

Background: "invalid argument", "unknown mode", "not supported": invalid enum-like argument errors explained — this error's family across 19 libraries.

Related errors

• top_k is not compatible with complex inputs.
• Arguments to batch_matmul must be at least 2D, got
• Arguments to batch_matmul must have same ndim, got
• axis argument out of range
• bfloat16 top_k is not supported on TPUv5 or older

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

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