{"record":{"id":"f8b8d6623b6a55f6","repo":"jax-ml/jax","slug":"dtypes-of-sorted-arr-and-query-must-match-got-so","errorCode":null,"errorMessage":"dtypes of sorted_arr and query must match; got {sorted_arr_aval.dtype} and {query_aval.dtype}","messagePattern":"dtypes of sorted_arr and query must match; got (.+?) and (.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/numpy/hijax.py","lineNumber":79,"sourceCode":"      *,\n      side: str,\n      dimension: int,\n      batch_dims: int,\n      method: str,\n      out_dtype: np.dtype):\n    batch_dims = operator.index(batch_dims)\n    if batch_dims < 0 or batch_dims >= sorted_arr_aval.ndim:\n      raise ValueError(\n          f\"batch_dims={batch_dims} must be in range [0, {sorted_arr_aval.ndim})\"\n      )\n    dimension = operator.index(dimension)\n    if not batch_dims <= dimension < sorted_arr_aval.ndim:\n      raise ValueError(\n          f\"dimension={dimension} must be in range [{batch_dims},\"\n          f\" {sorted_arr_aval.ndim})\"\n      )\n    if sorted_arr_aval.dtype != query_aval.dtype:\n      raise ValueError(\n          \"dtypes of sorted_arr and query must match; got \"\n          f\"{sorted_arr_aval.dtype} and {query_aval.dtype}\"\n      )\n    if side not in [\"left\", \"right\"]:\n      raise ValueError(\n          f\"invalid argument side={side!r}, expected 'left' or 'right'\"\n      )\n    if method not in self.valid_methods:\n      raise ValueError(\n          f\"invalid argument {method=}, expected one of {list(self.valid_methods)}\"\n      )\n    if sorted_arr_aval.shape[:batch_dims] != query_aval.shape[:batch_dims]:\n      raise ValueError(\n          \"batch dimension sizes must match; got\"\n          f\" {sorted_arr_aval.shape[:batch_dims]} != {query_aval.shape[:batch_dims]}\"\n      )\n    if not dtypes.issubdtype(out_dtype, np.integer):\n      raise ValueError(f\"out_dtype should be an integer type; got {out_dtype}\")","sourceCodeStart":61,"sourceCodeEnd":97,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/numpy/hijax.py#L61-L97","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","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"],"exampleFix":"// before\nidx = jnp.searchsorted(sorted_f32, query_f64)\n// after\nidx = jnp.searchsorted(sorted_f32, query_f64.astype(sorted_f32.dtype))","handlingStrategy":"validation","validationCode":"assert jnp.asarray(sorted_arr).dtype == jnp.asarray(query).dtype, (sorted_arr.dtype, query.dtype)","typeGuard":"def same_dtype(a, b) -> bool:\n    return jnp.asarray(a).dtype == jnp.asarray(b).dtype","tryCatchPattern":"try:\n    idx = jnp.searchsorted(a, q)\nexcept ValueError as e:\n    if 'dtypes of sorted_arr and query' in str(e):\n        q = q.astype(a.dtype)\n        idx = jnp.searchsorted(a, q)\n    else:\n        raise","preventionTips":["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"],"tags":["jax","dtype","searchsorted","validation"],"backgroundTag":"dtype-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}