code: def _get_hlo(f):
      def wrapped(*args, **kwargs):
          c = jax.xla_computation(f)(*args, **kwargs)
          print_opts = xla_client._xla.HloPrintOptions.short_parsable()
          print_opts.print_metadata = True
          return c.as_hlo_module().to_string(print_opts)
  
      return wrapped
  
imports: import functools
  from absl.testing import absltest
  import jax
  import jax.numpy
  from jax._src import core
  from jax import lax
  from jax._src.pjit import pjit
  from jax._src import linear_util
  from jax import config
  from jax._src import test_util
  from jax._src.lib import xla_client
label_desc: Use image node_id 1 for calling a global function with example usage: _get_hlo(f) and returns: wrapped, c