{"record":{"id":"7e8b0ea9d9bc3ef7","repo":"jax-ml/jax","slug":"logic-for-nan-policy-of-nan-policy-is-not-impl","errorCode":null,"errorMessage":"Logic for `nan_policy` of {nan_policy} is not implemented","messagePattern":"Logic for `nan_policy` of (.+?) is not implemented","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/scipy/stats/_core.py","lineNumber":97,"sourceCode":"    >>> mode, count = jax.scipy.stats.mode(x1, axis=1, keepdims=True)\n    >>> mode, count\n    (Array([[1],\n           [3],\n           [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:","sourceCodeStart":79,"sourceCodeEnd":115,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/scipy/stats/_core.py#L79-L115","documentation":"jax.scipy.stats.mode accepts nan_policy='omit' in its signature check but has not implemented the omit logic, so it raises NotImplementedError with a TODO. Only nan_policy='propagate' is actually functional in JAX.","triggerScenarios":"Calling jax.scipy.stats.mode(a, nan_policy='omit').","commonSituations":"Porting scipy.stats.mode code that used 'omit' to skip NaNs in survey/sensor data; expecting parity with scipy.","solutions":["Filter NaNs out of the array yourself before calling mode (e.g. x[~jnp.isnan(x)] for 1-D input).","Call scipy.stats.mode on the host with nan_policy='omit' instead.","File/check the upstream JAX issue and fall back to a manual nan-aware mode computation."],"exampleFix":"// before\njax.scipy.stats.mode(x, nan_policy='omit')\n\n// after\nclean = x[~jnp.isnan(x)]\njax.scipy.stats.mode(clean, nan_policy='propagate')","handlingStrategy":"fallback","validationCode":"# pre-clean NaNs instead of nan_policy='omit'\nx = x[~jnp.isnan(x)] if x.ndim == 1 else x","typeGuard":null,"tryCatchPattern":"try:\n    jax.scipy.stats.mode(x, nan_policy='omit')\nexcept NotImplementedError:\n    clean = x[~jnp.isnan(x)]\n    jax.scipy.stats.mode(clean)","preventionTips":["Strip NaNs before jax scipy.stats calls.","Track JAX scipy parity gaps when porting code."],"tags":["jax","scipy","stats","mode","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"}