jax-ml/jax · error · ValueError

Unsupported method: {method}

Error message

Unsupported method: {method}

What it means

Raised inside _searchsorted_impl when the method string is none of 'scan', 'scan_unrolled', 'compare_all', 'sort'. This is a defensive check behind the SearchSorted primitive: normally the constructor (error 1402) already validated method, so seeing this means the impl was called directly (e.g. via expand or eval_shape) with an invalid method.

Source

Thrown at jax/_src/numpy/hijax.py:547

    return tuple(jnp.where(fill_mask, 0, entry) for entry in out)


def _searchsorted_impl(sorted_arr: ArrayLike, query: ArrayLike, *, dimension: int,
                       batch_dims: int, side: str, dtype: np.dtype, method: str):
  """Main implementation of searchsorted primitive."""
  sorted_arr = jnp.moveaxis(sorted_arr, dimension, -1)
  dtype = dtypes._maybe_canonicalize_explicit_dtype(dtype, "searchsorted")

  if method == "scan":
    impl: Callable[..., Array] = functools.partial(_searchsorted_scan_impl, unrolled=False)
  elif method == "scan_unrolled":
    impl = functools.partial(_searchsorted_scan_impl, unrolled=True)
  elif method == "compare_all":
    impl = _searchsorted_compare_all_impl
  elif method == "sort":
    impl = _searchsorted_sort_impl
  else:
    raise ValueError(f"Unsupported method: {method}")

  fun = functools.partial(impl, side=side, dtype=dtype)
  for _ in range(sorted_arr.ndim - batch_dims - 1):
    fun = api.vmap(fun, in_axes=(0, None))
  for _ in range(batch_dims):
    fun = api.vmap(fun, in_axes=0)
  return fun(sorted_arr, query)


@api.jit(static_argnames=["side", "dtype", "unrolled"])
def _searchsorted_scan_impl(
    sorted_arr: Array, query: Array, side: str, dtype: np.dtype, unrolled: bool
) -> Array:
  """Scan-based implementation of searchsorted."""
  assert sorted_arr.ndim == 1
  assert side in ["left", "right"]
  (n,) = sorted_arr.shape
  if sorted_arr.size == 0:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Route through the public searchsorted primitive so constructor validation catches bad methods with a clearer message
  2. Use one of 'compare_all', 'scan', 'scan_unrolled', 'sort'

Example fix

# before
fun = functools.partial(_searchsorted_impl, method='binary', ...)
# after
fun = functools.partial(_searchsorted_impl, method='scan', ...)
# or better: use the public searchsorted API
Defensive patterns

Strategy: validation

Validate before calling

assert method in ('compare_all', 'scan', 'scan_unrolled', 'sort'), method

Type guard

def valid_method(m: str) -> bool:
    return m in ('compare_all', 'scan', 'scan_unrolled', 'sort')

Prevention

When it happens

Trigger: Calling the private _searchsorted_impl with method='binary' or another typo'd name directly, or bypassing the primitive's constructor validation when re-creating it.

Common situations: Library code calling the private impl directly (functools.partial(_searchsorted_impl, method=...)) instead of going through the public primitive; stale method strings from an older API version.

Related errors


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