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
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`.
  
tensorflow_transform/py_func/pyfunc_helper.py