code: def median(
      input: JaxArray,
      /,
      *,
      axis: Optional[Union[Tuple[int], int]] = None,
      keepdims: bool = False,
      out: Optional[JaxArray] = None,
  ) -> JaxArray:
      if isinstance(axis, list):
          axis = tuple(axis)
      ret = jnp.median(
          input,
          axis=axis,
          keepdims=keepdims,
          out=out,
      )
      if input.dtype in [jnp.uint64, jnp.int64, jnp.float64]:
          return ret.astype(jnp.float64)
      elif input.dtype in [jnp.float16, jnp.bfloat16]:
          return ret.astype(input.dtype)
      else:
          return ret.astype(jnp.float32)
  
imports: import jax.numpy
  from typing import Optional, Union, Tuple, Sequence
  from ivy.functional.backends.jax import JaxArray
  import jax.lax
  import ivy
  from ivy.func_wrapper import with_unsupported_dtypes
  from .None import backend_version
  from ..statistical import _infer_dtype
label_desc: Use image node_id 2 for calling a global function with example usage: median() and returns: ret, ret, ret