code: def tf_wrap_with_input_names(f, input_shapes):
      def wrapper(*args):
          args = tuple(
              tf.identity(a, name=name)
              for a, (name, _) in zip(args, input_shapes)
          )
          # NOTE: Output names already set via `jax2tf.convert(..)`.
          return f(*args)
  
      return wrapper
  
file_path: jax/jax/tools/jax_to_ir.py
imports: from ast import literal_eval
  import importlib
  import functools
  import re
  from absl import app
  from absl import flags
  import jax
  import jax.numpy
label_desc: Use image node_id 2 for calling a global function with example usage: tf_wrap_with_input_names(f, input_shapes) and returns: wrapper, f