jax-ml/jax · error · ValueError

invalid argument {method=}, expected one of {list(self.valid

Error message

invalid argument {method=}, expected one of {list(self.valid_methods)}

What it means

Raised by the SearchSorted HiJAX primitive when the method argument is not one of ('compare_all', 'scan', 'scan_unrolled', 'sort'). method selects the lower-level implementation strategy for the search (e.g. binary scan, full sort, or pairwise compare) used inside the primitive.

Source

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

          f"batch_dims={batch_dims} must be in range [0, {sorted_arr_aval.ndim})"
      )
    dimension = operator.index(dimension)
    if not batch_dims <= dimension < sorted_arr_aval.ndim:
      raise ValueError(
          f"dimension={dimension} must be in range [{batch_dims},"
          f" {sorted_arr_aval.ndim})"
      )
    if sorted_arr_aval.dtype != query_aval.dtype:
      raise ValueError(
          "dtypes of sorted_arr and query must match; got "
          f"{sorted_arr_aval.dtype} and {query_aval.dtype}"
      )
    if side not in ["left", "right"]:
      raise ValueError(
          f"invalid argument side={side!r}, expected 'left' or 'right'"
      )
    if method not in self.valid_methods:
      raise ValueError(
          f"invalid argument {method=}, expected one of {list(self.valid_methods)}"
      )
    if sorted_arr_aval.shape[:batch_dims] != query_aval.shape[:batch_dims]:
      raise ValueError(
          "batch dimension sizes must match; got"
          f" {sorted_arr_aval.shape[:batch_dims]} != {query_aval.shape[:batch_dims]}"
      )
    if not dtypes.issubdtype(out_dtype, np.integer):
      raise ValueError(f"out_dtype should be an integer type; got {out_dtype}")
    # Attempt this here to catch overflow errors early.
    out_dtype.type(sorted_arr_aval.shape[dimension])
    self.in_avals = (sorted_arr_aval, query_aval)
    self.out_aval = core.typeof(api.eval_shape(
      functools.partial(_searchsorted_impl,
        dimension=dimension, batch_dims=batch_dims, side=side,
        dtype=out_dtype, method=method),
        sorted_arr_aval, query_aval))
    self.params = dict(

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Use one of the methods printed in the error: 'compare_all', 'scan', 'scan_unrolled', or 'sort'
Leave method at its default rather than guessing names

Example fix

# before
idx = jax_hijax_searchsorted(a, v, method='binary')
# after
idx = jax_hijax_searchsorted(a, v, method='scan')

Defensive patterns

Strategy: validation

Validate before calling

VALID = ('compare_all', 'scan', 'scan_unrolled', 'sort')
assert method in VALID, f'{method=} not in {VALID}'

Type guard

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

Prevention

Leave method at default unless benchmarking
Treat method names as an enum; add a unit test enumerating valid values

When it happens

Trigger: Calling searchsorted with method='binary' (the intuitive but invalid guess), method='scans', or any string outside the four valid methods listed in the error message.

Common situations: Assuming numpy-compatible parameter names; copy-pasting method names from a different JAX version or a different search API.

Related errors

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