code: def register_pyfuncs_from_saved_transform(
      graph, meta_graph, loaded_in_tf2
  ):
      """Registers `py_func`s in the MetaGraphDef.
  
      Takes the picked `py_func`s stored in the MetaGraphDef and adds them to the
      graph.  Registered `py_func`s are referred to internally by the token
      attribute of the `py_func` op.  We first create some arbitrary ops which
      are not used, but which result in the pickled functions stored in the
      MetaGraphDef being registered.  We then take the tokens of these newly
      registered functions, and remap the tokens in the MetaGraphDef to contain
      the new tokens for each function (this remapping is required since we cannot
      specify what token should be used to register a function).
  
      Args:
        graph: The tf.Graph into which the meta_graph_def will be imported.
        meta_graph: The MetaGraphDef containing the `py_func`s.  All the `py_func`
           ops in the graph will be modified in-place to have their token point to
           the newly regsitered function.
        loaded_in_tf2: A boolean indicating whether the saved transform is being
           re-loaded in TF1 or TF2.
  
      Returns:
        Modified graph_def if pyfuncs were found, else None.
  
      Raises:
        ValueError if an unregistered pyfunc is encountered in `graph`.
      """
      if _PYFUNC_COLLECTION_KEY not in meta_graph.collection_def:
          return None
  
      # TODO(b/35929054) to enable it in TF itself. Once supported,
      # we should refactor this code to remove extra work for pickling and
      # re-registering of the py_funcs.
      pyfuncs_collection = meta_graph.collection_def[
          _PYFUNC_COLLECTION_KEY
      ]
  
      new_tokens_by_old_token = {}
      with graph.as_default():
          for func_def_str in pyfuncs_collection.bytes_list.value:
              func_def = _PyFuncDef.from_proto_string(func_def_str)
              # Re-insert the original python function into the default graph.
              # The operation itself in the graph does not matter (hence the dummy
              # values for name, Tout, and stateful). This is done only to reinsert
              # the function body in the internal TF's function registry.
              # TODO(b/123241062): We should even remove this op from the graph if
              # possible.
              func_temp_name = func_def.token + b"_temp"
              output_tensor = insert_pyfunc(
                  func_def.func, tf.float32, False, func_temp_name
              )
              # Store the token associated with the function associated with the call
              # to tf.py_func.
              token = output_tensor.op.get_attr("token")
              new_tokens_by_old_token[func_def.token] = token
  
      if loaded_in_tf2:
          graph_def = graph.as_graph_def()
          # Since we are updating the GraphDef of the graph in whose context pyfuncs
          # were re-inserted, new tokens will also be present.
          expected_tokens_in_graph_def = list(
              new_tokens_by_old_token.keys()
          ) + list(new_tokens_by_old_token.values())
      else:
          graph_def = meta_graph.graph_def
          expected_tokens_in_graph_def = new_tokens_by_old_token.keys()
      # Swap the old token stored for the function with the new one, if there are
      # any tokens to change.
      if new_tokens_by_old_token:
          for node in graph_def.node:
              if node.op == "PyFunc" or node.op == "PyFuncStateless":
                  token = node.attr["token"]
                  new_token = new_tokens_by_old_token.get(token.s, None)
                  if new_token is not None:
                      token.s = new_token
                  else:
                      if token.s not in expected_tokens_in_graph_def:
                          raise ValueError(
                              f"Function: {node.name} was not registered"
                          )
      return graph_def
  
file_path: transform/tensorflow_transform/py_func/pyfunc_helper.py
imports: import dill
  import tensorflow
  from tfx_bsl.types import tfx_namedtuple
  from tensorflow.core.framework import attr_value_pb2
  from tensorflow.python.framework import ops
label_desc: Use image node_id 2 for calling a global function with example usage: register_pyfuncs_from_saved_transform(graph, meta_graph, loaded_in_tf2) and returns: graph_def, None