{"record":{"id":"bdd6df3e77cb15ad","repo":"jax-ml/jax","slug":"in-order-to-best-jit-compile-rankdata-we-cannot","errorCode":null,"errorMessage":"In order to best JIT compile `rankdata`, we cannot know whether `x` contains nans. Please check if nans exist in `x` outside of the `rankdata` function.","messagePattern":"In order to best JIT compile `rankdata`, we cannot know whether `x` contains nans\\. Please check if nans exist in `x` outside of the `rankdata` function\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/scipy/stats/_core.py","lineNumber":190,"sourceCode":"    Array([1., 3., 2.], dtype=float32)\n\n    >>> x = jnp.array([1, 3, 2, 3])\n    >>> rankdata(x)\n    Array([1. , 3.5, 2. , 3.5], dtype=float32)\n  \"\"\"\n  check_arraylike(\"rankdata\", a)\n\n  if nan_policy not in [\"propagate\", \"omit\", \"raise\"]:\n    raise ValueError(\n      f\"Illegal nan_policy value {nan_policy!r}; expected one of \"\n      \"{'propagate', 'omit', 'raise'}\"\n    )\n  if nan_policy == \"omit\":\n    raise NotImplementedError(\n      f\"Logic for `nan_policy` of {nan_policy} is not implemented\"\n    )\n  if nan_policy == \"raise\":\n    raise NotImplementedError(\n      \"In order to best JIT compile `rankdata`, we cannot know whether `x` \"\n      \"contains nans. Please check if nans exist in `x` outside of the \"\n      \"`rankdata` function.\"\n    )\n\n  if method not in (\"average\", \"min\", \"max\", \"dense\", \"ordinal\"):\n    raise ValueError(f\"unknown method '{method}'\")\n\n  if axis is not None:\n    return jnp.apply_along_axis(rankdata, axis, a, method)\n\n  a = jnp.ravel(a)\n  out_dtype = dtypes.default_float_dtype()\n\n  def _rankdata(a: Array) -> Array:\n    arr, sorter = lax.sort_key_val(a, jnp.arange(a.size))\n    inv = invert_permutation(sorter)\n","sourceCodeStart":172,"sourceCodeEnd":208,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/scipy/stats/_core.py#L172-L208","documentation":"jax.scipy.stats.rankdata with nan_policy='raise' raises NotImplementedError because NaN detection is data-dependent and cannot be JIT-traced; users must check for NaNs outside the function.","triggerScenarios":"Calling jax.scipy.stats.rankdata(a, nan_policy='raise') — unconditionally raises, NaNs present or not.","commonSituations":"Wanting fail-fast NaN semantics inside jitted pipelines; ported scipy code with 'raise' policy.","solutions":["Check for NaNs before calling: if jnp.isnan(x).any(): raise ...; then use nan_policy='propagate' or omit the kwarg.","Move the NaN check to the data-preparation stage of the pipeline."],"exampleFix":"// before\njax.scipy.stats.rankdata(x, nan_policy='raise')\n\n// after\nassert not jnp.isnan(x).any(), \"x contains NaNs\"\njax.scipy.stats.rankdata(x)  # default: propagate","handlingStrategy":"validation","validationCode":"if jnp.isnan(x).any():\n    raise ValueError(\"x contains NaNs\")\njax.scipy.stats.rankdata(x)","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Fail fast on NaNs in data prep, not inside stats calls.","Never pass nan_policy='raise' to jax scipy.stats."],"tags":["jax","scipy","stats","rankdata","jit","not-implemented","nan-policy"],"backgroundTag":"scipy-nan-policy-unsupported","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}