code: def scalar_type_of(x: Any) -> type:
      """Return the scalar type associated with a JAX value."""
      typ = dtype(x)
      if typ in _custom_float_dtypes:
          return float
      elif typ in _int4_dtypes:
          return int
      elif np.issubdtype(typ, np.bool_):
          return bool
      elif np.issubdtype(typ, np.integer):
          return int
      elif np.issubdtype(typ, np.floating):
          return float
      elif np.issubdtype(typ, np.complexfloating):
          return complex
      else:
          raise TypeError(f"Invalid scalar value {x}")
  
imports: from __future__ import annotations
  import abc
  import builtins
  import functools
  from typing import cast, overload, Any, Literal, Union
  import warnings
  import ml_dtypes
  import numpy
  from jax._src import config
  from jax._src.typing import DType, DTypeLike
  from jax._src.util import set_module
  from jax._src import traceback_util
label_desc: Use image node_id 9 for calling a global function with example usage: scalar_type_of(x) and returns: float, int, bool, int, float, complex