{"record":{"id":"d2ddc2afd65ac0e3","repo":"jax-ml/jax","slug":"in-order-to-best-jit-compile-mode-we-cannot-kno","errorCode":null,"errorMessage":"In order to best JIT compile `mode`, we cannot know whether `x` contains nans. Please check if nans exist in `x` outside of the `mode` function.","messagePattern":"In order to best JIT compile `mode`, we cannot know whether `x` contains nans\\. Please check if nans exist in `x` outside of the `mode` function\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/scipy/stats/_core.py","lineNumber":101,"sourceCode":"           [2]], dtype=int32), Array([[3],\n           [3],\n           [3]], dtype=int32))\n  \"\"\"\n  check_arraylike(\"mode\", a)\n  x = jnp.atleast_1d(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    # TODO: return answer without nans included.\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 `mode`, we cannot know whether `x` contains nans. \"\n      \"Please check if nans exist in `x` outside of the `mode` function.\"\n    )\n  if axis is not None:\n    axis = canonicalize_axis(axis, x.ndim)\n\n  input_shape = x.shape\n  if keepdims:\n    if axis is None:\n      output_shape = tuple(1 for i in input_shape)\n    else:\n      output_shape = tuple(1 if i == axis else s for i, s in enumerate(input_shape))\n  else:\n    if axis is None:\n      output_shape = ()\n    else:\n      output_shape = tuple(s for i, s in enumerate(input_shape) if i != axis)\n","sourceCodeStart":83,"sourceCodeEnd":119,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/scipy/stats/_core.py#L83-L119","documentation":"jax.scipy.stats.mode with nan_policy='raise' raises NotImplementedError because detecting NaNs would require a data-dependent branch, which breaks JIT tracing. The message asks users to check for NaNs outside the function.","triggerScenarios":"Calling jax.scipy.stats.mode(a, nan_policy='raise') — even when the array contains no NaNs, this always raises.","commonSituations":"Porting scipy defaults (scipy used 'propagate' historically but users often set 'raise'); wrapping mode in jit and wanting NaN safety.","solutions":["Remove nan_policy='raise' and check for NaNs before the call: has_nan = bool(jnp.isnan(x).any()); raise manually if true.","Validate input cleanliness upstream (assert not jnp.isnan(x).any()) and then call mode with 'propagate'.","Use scipy.stats.mode outside JIT when 'raise' semantics are needed."],"exampleFix":"// before\njax.scipy.stats.mode(x, nan_policy='raise')\n\n// after\nif jnp.isnan(x).any():\n    raise ValueError(\"x contains NaNs\")\njax.scipy.stats.mode(x, nan_policy='propagate')","handlingStrategy":"validation","validationCode":"if jnp.isnan(x).any():\n    raise ValueError(\"x contains NaNs\")\njax.scipy.stats.mode(x, nan_policy='propagate')","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Do NaN checks eagerly outside jitted functions.","Avoid nan_policy='raise' with jax; it always raises."],"tags":["jax","scipy","stats","mode","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"}