jax-ml/jax · error · ValueError

Invalid mode '{mode}' for np.take

Error message

Invalid mode '{mode}' for np.take

What it means

jnp.take accepts only modes None/'fill', 'raise' (unimplemented), 'wrap', and 'clip'. Any other string (e.g. 'promiselike', typo) triggers this ValueError.

Source

Thrown at jax/_src/numpy/indexing.py:726

    a = a.ravel()
    axis_idx = 0
  else:
    axis_idx = canonicalize_axis(axis, np.ndim(a))

  if mode is None or mode == "fill":
    gather_mode = slicing.GatherScatterMode.FILL_OR_DROP
    # lax.gather() does not support negative indices, so we wrap them here
    indices = util._where(indices < 0, indices + a.shape[axis_idx], indices)
  elif mode == "raise":
    # TODO(phawkins): we have no way to report out of bounds errors yet.
    raise NotImplementedError("The 'raise' mode to jnp.take is not supported.")
  elif mode == "wrap":
    indices = ufuncs.mod(indices, lax._const(indices, a.shape[axis_idx]))
    gather_mode = slicing.GatherScatterMode.PROMISE_IN_BOUNDS
  elif mode == "clip":
    gather_mode = slicing.GatherScatterMode.CLIP
  else:
    raise ValueError(f"Invalid mode '{mode}' for np.take")

  index_dims = len(np.shape(indices))
  slice_sizes = list(np.shape(a))
  if slice_sizes[axis_idx] == 0:
    if indices.size != 0:
      raise IndexError("Cannot do a non-empty jnp.take() from an empty axis.")
    return a

  if indices.size == 0:
    out_shape = (slice_sizes[:axis_idx] + list(indices.shape) +
                 slice_sizes[axis_idx + 1:])
    return lax.full_like(a, 0, shape=out_shape)

  slice_sizes[axis_idx] = 1
  dnums = slicing.GatherDimensionNumbers(
    offset_dims=tuple(
      list(range(axis_idx)) +
      list(range(axis_idx + index_dims, len(a.shape) + index_dims - 1))),

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Use one of: None, 'fill', 'clip', 'wrap'
  2. Check spelling of the mode string against the docs

Example fix

// before
y = jnp.take(a, idx, mode='fills')
// after
y = jnp.take(a, idx, mode='fill')
Defensive patterns

Strategy: validation

Validate before calling

assert mode in (None, 'fill', 'clip', 'wrap'), f'invalid take mode {mode!r}'

Prevention

When it happens

Trigger: Calling jnp.take(..., mode='weap') or any unrecognized mode string; passing a mode variable that's None-adjacent or misspelled.

Common situations: Typos in mode strings; code parameterizing mode from config where invalid values slip through.

Understand the failure class

Background: "Must be a positive integer", "Invalid value", "Unsupported": the invalid-argument-value error family, when a library rejects the value you pass — this error's family across 35 libraries.

Related errors


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