code: def jit(f, device=None):
      if device is None:
          device = utils.get_first_accelerator()
      # TODO(tomhennigan) Enable XLA compilation (experimental_compile=True).
      return tf.function(utils.run_on_device(f, device))
  
file_path: sonnet/sonnet/src/functional/jax.py
imports: import functools
  from sonnet.src.functional import utils
  import tensorflow
  import tree
label_desc: Use image node_id 3 for calling a global function with example usage: jit(f, device) and returns: tf