jax-ml/jax · error · ValueError

dtypes of sorted_arr and query must match; got {sorted_arr_a

Error message

dtypes of sorted_arr and query must match; got {sorted_arr_aval.dtype} and {query_aval.dtype}

What it means

Raised by the SearchSorted HiJAX primitive constructor when the sorted array and the query array have different dtypes. jax.numpy.searchsorted (the HiJAX implementation) requires both operands to share the same dtype because the comparison kernels are dtype-specialized. The error is raised eagerly at primitive-construction time, before any computation runs.

Source

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

      *,
      side: str,
      dimension: int,
      batch_dims: int,
      method: str,
      out_dtype: np.dtype):
    batch_dims = operator.index(batch_dims)
    if batch_dims < 0 or batch_dims >= sorted_arr_aval.ndim:
      raise ValueError(
          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}")

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Cast query to the sorted array's dtype: query = query.astype(sorted_arr.dtype)
Or cast both to a common dtype: jax.numpy.searchsorted(a.astype(jnp.float32), q.astype(jnp.float32))
Check jax.config.x64_enabled if you intended 64-bit dtypes and they silently became float32/int32

Example fix

// before
idx = jnp.searchsorted(sorted_f32, query_f64)
// after
idx = jnp.searchsorted(sorted_f32, query_f64.astype(sorted_f32.dtype))

Defensive patterns

Strategy: validation

Validate before calling

assert jnp.asarray(sorted_arr).dtype == jnp.asarray(query).dtype, (sorted_arr.dtype, query.dtype)

Type guard

def same_dtype(a, b) -> bool:
    return jnp.asarray(a).dtype == jnp.asarray(b).dtype

Try / catch

try:
    idx = jnp.searchsorted(a, q)
except ValueError as e:
    if 'dtypes of sorted_arr and query' in str(e):
        q = q.astype(a.dtype)
        idx = jnp.searchsorted(a, q)
    else:
        raise

Prevention

Canonicalize dtypes at data-load boundaries with astype
Standardize one float precision (e.g. float32) across the pipeline
Watch x64 config changes that silently alter default dtypes

When it happens

Trigger: Calling jax.numpy searchsorted where sorted_arr and query have different dtypes, e.g. a float32 sorted array with a float64 query, or int32 vs int32/weak-typed Python scalars that canonicalize differently under JAX's default dtype promotion.

Common situations: Loading data with np.loadtxt/np.genfromtxt (float64) and searching against jnp.float32 arrays; mixing arrays created under different x64 enabled/disabled configurations; passing Python floats that get weak-typed to a different default than the array.

Related errors

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