What causes "The 'raise' mode to jnp.take is not supported." in jax-ml/jax?

Use mode='clip' or mode='wrap' (or default fill behavior) instead

jax-ml/jax · error · NotImplementedError

The 'raise' mode to jnp.take is not supported.

Error message

The 'raise' mode to jnp.take is not supported.

What it means

NumPy's mode='raise' for np.take raises on out-of-bound indices at runtime; JAX cannot support this because errors cannot be raised from jit-compiled/XLA code, so it raises NotImplementedError up front.

Source

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

def _take(a, indices, axis: int | None = None, out=None, mode=None,
          unique_indices=False, indices_are_sorted=False, fill_value=None):
  if out is not None:
    raise NotImplementedError("The 'out' argument to jnp.take is not supported.")
  a, indices = util.ensure_arraylike("take", a, indices)

  if axis is None:
    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:])

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Use mode='clip' or mode='wrap' (or default fill behavior) instead
Validate indices in Python before the call: assert ((indices >= 0) & (indices < n)).all() when indices are concrete

Example fix

// before
y = jnp.take(a, idx, mode='raise')
// after
y = jnp.take(a, idx, mode='clip')  # or validate idx beforehand

Defensive patterns

Strategy: fallback

Validate before calling

assert mode != 'raise', "mode='raise' unsupported; use 'clip' or 'wrap'"

Prevention

Replace mode='raise' with clip/wrap plus explicit validation
Validate concrete indices in Python: ((idx >= 0) & (idx < n)).all()

When it happens

Trigger: Calling jnp.take(a, indices, mode='raise').

Common situations: Ported NumPy code that relied on mode='raise' for bounds checking; debugging index computations in NumPy before moving to JAX.

Related errors

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