code: def cast_outputs(fun):
      def wrapped(*args, **kwargs):
          dtype = np.asarray(args[0]).dtype
          return jax.tree_map(
              lambda x: np.asarray(x, dtype=dtype), fun(*args, **kwargs)
          )
  
      return wrapped
  
imports: from functools import partial
  from absl.testing import absltest
  import numpy
  import jax
  import jax.numpy
  from jax._src import test_util
  from jax._src.numpy.ufunc_api import get_if_single_primitive
  from jax import config
label_desc: Use image node_id 5 for calling a global function with example usage: cast_outputs(fun) and returns: wrapped, jax