code: def _jax_type(dtype: DType, weak_type: bool) -> JAXType:
      """Return the jax type for a dtype and weak type."""
      if weak_type:
          if dtype == bool:
              return dtype
          if dtype in _custom_float_dtypes:
              return float
          return type(dtype.type(0).item())
      return dtype
  
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 14 for calling a global function with example usage: _jax_type(dtype, weak_type) and returns: dtype, type, dtype, float