jax-ml/jax · error · NotImplementedError

jax.scipy.ndimage.map_coordinates does not yet support mode

Error message

jax.scipy.ndimage.map_coordinates does not yet support mode {}. Currently supported modes are {}.

What it means

JAX's map_coordinates only implements a subset of scipy's boundary modes (those in _INDEX_FIXERS). Unsupported modes like 'wrap', 'reflect' (older versions), 'grid-wrap', or 'mirror' raise NotImplementedError.

Source

Thrown at jax/_src/scipy/ndimage.py:85

  lower_weight = 1 - upper_weight
  index = lower.astype(np.int32)
  return [(index, lower_weight), (index + 1, upper_weight)]


@api.jit(static_argnums=(2, 3, 4))
def _map_coordinates(input: ArrayLike, coordinates: Sequence[ArrayLike],
                     order: int, mode: str, cval: ArrayLike) -> Array:
  input_arr = jnp.asarray(input)
  coordinate_arrs = [jnp.asarray(c) for c in coordinates]
  cval = jnp.asarray(cval, input_arr.dtype)

  if len(coordinates) != input_arr.ndim:
    raise ValueError('coordinates must be a sequence of length input.ndim, but '
                     '{} != {}'.format(len(coordinates), input_arr.ndim))

  index_fixer = _INDEX_FIXERS.get(mode)
  if index_fixer is None:
    raise NotImplementedError(
        'jax.scipy.ndimage.map_coordinates does not yet support mode {}. '
        'Currently supported modes are {}.'.format(mode, set(_INDEX_FIXERS)))

  if mode == 'constant':
    is_valid = lambda index, size: (0 <= index) & (index < size)
  else:
    is_valid = lambda index, size: True

  if order == 0:
    interp_fun = _nearest_indices_and_weights
  elif order == 1:
    interp_fun = _linear_indices_and_weights
  else:
    raise NotImplementedError(
        'jax.scipy.ndimage.map_coordinates currently requires order<=1')

  valid_1d_interpolations = []
  for coordinate, size in zip(coordinate_arrs, input_arr.shape):

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use a supported mode such as 'constant', 'nearest', or the reflect variants listed in the error message
2. Upgrade jax/jaxlib — more modes have been added over time
3. Implement the boundary fix-up yourself by clamping coordinates before the call

Example fix

# before
out = ndimage.map_coordinates(img, coords, mode='wrap', order=1)
# after
out = ndimage.map_coordinates(img, coords, mode='nearest', order=1)

from jax._src.scipy.ndimage import _INDEX_FIXERS
mode = mode if mode in _INDEX_FIXERS else 'nearest'

def mode_supported(mode: str) -> bool: return mode in {'constant','nearest','reflect','mirror'}  # verify for your jax version

try:
    out = map_coordinates(img, coords, mode=mode)
except NotImplementedError:
    out = map_coordinates(img, coords, mode='nearest')

When it happens

Trigger: Passing mode='wrap' or mode='mirror' on a JAX version that has not implemented it; copying scipy.ndimage.map_coordinates calls verbatim.

Common situations: Porting scipy image pipelines (e.g. elastic deformations) to JAX and hitting missing mode parity.

Related errors

• jax.scipy.ndimage.map_coordinates currently requires order<=
• coordinates must be a sequence of length input.ndim, but {}
• ind must be a positive integer; got {ind=}
• tensorinv is only possible when the product of the first `in
• collective_axes is not supported in pallas_call. Use plgpu.k

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