jax-ml/jax · error · NotImplementedError

jax.scipy.ndimage.map_coordinates currently requires order<=

Error message

jax.scipy.ndimage.map_coordinates currently requires order<=1

What it means

JAX's map_coordinates only implements order 0 (nearest) and order 1 (linear) interpolation; any higher spline order raises NotImplementedError because spline prefilters are not implemented.

Source

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

                     '{} != {}'.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):
    interp_nodes = interp_fun(coordinate)
    valid_interp = []
    for index, weight in interp_nodes:
      fixed_index = index_fixer(index, size)
      valid = is_valid(index, size)
      valid_interp.append((fixed_index, valid, weight))
    valid_1d_interpolations.append(valid_interp)

  outputs = []
  for items in itertools.product(*valid_1d_interpolations):
    indices, validities, weights = util.unzip3(items)
    if all(valid is True for valid in validities):
      # fast path
      contribution = input_arr[indices]

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Set order=0 or order=1
2. Accept the interpolation quality change, or pre/post-process with a hand-written higher-order interpolation
3. Watch jax releases; higher orders may land later

Example fix

# before
out = ndimage.map_coordinates(img, coords, order=3)
# after
out = ndimage.map_coordinates(img, coords, order=1)

order = 1 if order > 1 else order

def order_supported(o: int) -> bool: return o in (0, 1)

try:
    out = map_coordinates(img, coords, order=order)
except NotImplementedError:
    out = map_coordinates(img, coords, order=1)

When it happens

Trigger: Passing order=3 (scipy's default) or order=2..5; migrating scipy.ndimage.shift/rotate/affine_transform code that defaults to cubic splines.

Common situations: Porting image augmentation or registration pipelines that rely on cubic interpolation; accuracy differences after downgrading to linear.

Related errors

• jax.scipy.ndimage.map_coordinates does not yet support mode
• 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/36f72910cd25f6bb. Report an issue: GitHub.