jax-ml/jax · error · ValueError

out_dtype should be an integer type; got {out_dtype}

Error message

out_dtype should be an integer type; got {out_dtype}

What it means

Raised by the SearchSorted HiJAX primitive when the requested output dtype for the returned indices is not an integer type. Search results are indices, so out_dtype must satisfy dtypes.issubdtype(out_dtype, np.integer).

Source

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

      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(
      side=side,
      dimension=dimension,
      batch_dims=batch_dims,
      method=method,
    )
    super().__init__()

  def expand(self, sorted_arr: ArrayLike, query: ArrayLike) -> Array:  # pyrefly: ignore[bad-override]
    return _searchsorted_impl(

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Use an integer dtype such as 'int32' (the usual default) or 'int64'
If you need float positions, get integer indices first and cast the result afterwards

Example fix

# before
idx = searchsorted(a, v, out_dtype=jnp.float32)
# after
idx = searchsorted(a, v, out_dtype=jnp.int32).astype(jnp.float32)

Defensive patterns

Strategy: validation

Validate before calling

import numpy as np
from jax._src import dtypes
assert dtypes.issubdtype(np.dtype(out_dtype), np.integer)

Type guard

def is_int_dtype(d) -> bool:
    import numpy as np
    from jax._src import dtypes
    return dtypes.issubdtype(np.dtype(d), np.integer)

Prevention

Reserve dtype params of index-returning functions for int32/int64
Cast outputs after the call if floats are needed

When it happens

Trigger: Calling searchsorted with out_dtype='float32', np.float64, or a non-dtype object that resolves to a floating type.

Common situations: Copy-pasting a dtype parameter from a nearby array-creation call; attempting to get 'positions' as floats instead of indices.

Related errors

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