code: def _preserve_input_types(f):
      def wrapped(*args):
          dtype = _dtype(args[0])
          result = np.array(f(*args), dtype=dtype)
          if all(is_python_scalar(arg) for arg in args):
              result = result.item()
          return result
  
      return wrapped
  
imports: from functools import partial
  import operator
  from jax._src import api
  from jax._src import config
  from jax._src import dtypes
  from jax._src.tree_util import tree_map, tree_reduce
  import numpy
label_desc: Use image node_id 11 for calling a global function with example usage: _preserve_input_types(f) and returns: wrapped, result