code: def rankdata(
      a: ArrayLike,
      method: str = "average",
      *,
      axis: int | None = None,
      nan_policy: str = "propagate",
  ) -> Array:
      check_arraylike("rankdata", a)
  
      if nan_policy not in ["propagate", "omit", "raise"]:
          raise ValueError(
              f"Illegal nan_policy value {nan_policy!r}; expected one of "
              "{'propoagate', 'omit', 'raise'}"
          )
      if nan_policy == "omit":
          raise NotImplementedError(
              f"Logic for `nan_policy` of {nan_policy} is not implemented"
          )
      if nan_policy == "raise":
          raise NotImplementedError(
              "In order to best JIT compile `mode`, we cannot know whether `x` "
              "contains nans. Please check if nans exist in `x` outside of the "
              "`rankdata` function."
          )
  
      if method not in ("average", "min", "max", "dense", "ordinal"):
          raise ValueError(f"unknown method '{method}'")
  
      a = jnp.asarray(a)
  
      if axis is not None:
          return jnp.apply_along_axis(rankdata, axis, a, method)
  
      arr = jnp.ravel(a)
      sorter = jnp.argsort(arr)
      inv = invert_permutation(sorter)
  
      if method == "ordinal":
          return inv + 1
      arr = arr[sorter]
      obs = jnp.insert(arr[1:] != arr[:-1], 0, True)
      dense = obs.cumsum()[inv]
      if method == "dense":
          return dense
      count = jnp.nonzero(obs, size=arr.size + 1, fill_value=len(obs))[
          0
      ]
      if method == "max":
          return count[dense]
      if method == "min":
          return count[dense - 1] + 1
      if method == "average":
          return 0.5 * (count[dense] + count[dense - 1] + 1).astype(
              dtypes.canonicalize_dtype(jnp.float_)
          )
      raise ValueError(f"unknown method '{method}'")
  
file_path: jax/jax/_src/scipy/stats/_core.py
imports: from __future__ import annotations
  from collections import namedtuple
  from functools import partial
  import math
  import jax
  import jax.numpy
  from jax import jit
  from jax._src import dtypes
  from jax._src.api import vmap
  from jax._src.numpy.util import check_arraylike, _wraps, promote_args_inexact
  from jax._src.typing import ArrayLike, Array
  from jax._src.util import canonicalize_axis
  import scipy
label_desc: Use image node_id 3 for calling a global function with example usage: rankdata(a, method) and returns: jnp, unknown, dense, count, unknown, unknown