message: Refactor setup of revived functions from a SavedModel.
diff: @@ -56,24 +56,24 @@ class _Loader(object):
  function_deserialization.load_function_def_library(
  meta_graph.graph_def.library))
  self._load_all()
  - self._setup_functions()
  + # TODO(b/124045874): There are limitations with functions whose captures
  + # trigger other functions to be executed. For now it is only guaranteed to
  + # work if the captures of a function only trigger functions without
  + # captures.
  + self._setup_functions_structures()
  + self._setup_functions_captures()
  self._restore_checkpoint()
  - def _setup_concrete_function(self, proto, concrete_function, coder):
  - """Setup captured tensors and outputs for a single concrete function."""
  - bound_inputs = [
  - self._get_tensor_from_node(node_id)
  - for node_id in proto.bound_inputs]
  - bound_variables = [
  - self._nodes[node_id]
  - for node_id in proto.bound_inputs
  - if self._proto.nodes[node_id].WhichOneof("kind") == "variable"
  - ]
  - # TODO(andresp): This is only injecting the captured inputs into the
  - # concrete function, note that we did not modify the FuncGraph
  - # itself.
  - concrete_function._captured_inputs = bound_inputs # pylint: disable=protected-access
  - concrete_function._func_graph.variables = bound_variables # pylint: disable=protected-access
  + for node in self._nodes:
  + if isinstance(node, tracking.TrackableResource):
  + init_op = node.initialize()
  + ops.add_to_collection(ops.GraphKeys.TABLE_INITIALIZERS, init_op)
  +
  + def _setup_functions_structures(self):
  + """Setup structure for inputs and outputs of restored functions."""
  + coder = nested_structure_coder.StructureCoder()
  + for name, proto in sorted(self._proto.concrete_functions.items()):
  + concrete_function = self._concrete_functions[name]
  # By setting the structured_outputs directly, we can rely on this
  # function_lib.ConcreteFunction object to perform the output repacking
  # logic. The only limitation of that logic is that it only works
  @@ -91,14 +91,24 @@ class _Loader(object):
  concrete_function._func_graph.structured_input_signature = ( # pylint: disable=protected-access
  coder.decode_proto(proto.canonicalized_input_signature))
  - def _setup_functions(self):
  - """Setup captures and output structure in restored functions."""
  - coder = nested_structure_coder.StructureCoder()
  - for name, concrete_function_proto in self._proto.concrete_functions.items():
  - self._setup_concrete_function(
  - concrete_function_proto,
  - self._concrete_functions[name],
  - coder)
  + def _setup_functions_captures(self):
  + """Setup captures and variables in restored functions."""
  + concrete_functions = sorted(self._proto.concrete_functions.items())
  + for name, proto in concrete_functions:
  + concrete_function = self._concrete_functions[name]
  + bound_inputs = [
  + self._get_tensor_from_node(node_id)
  + for node_id in proto.bound_inputs]
  + bound_variables = [
  + self._nodes[node_id]
  + for node_id in proto.bound_inputs
  + if self._proto.nodes[node_id].WhichOneof("kind") == "variable"
  + ]
  + # TODO(andresp): This is only injecting the captured inputs into the
  + # concrete function, note that we did not modify the FuncGraph
  + # itself.
  + concrete_function._captured_inputs = bound_inputs # pylint: disable=protected-access
  + concrete_function._func_graph.variables = bound_variables # pylint: disable=protected-access
  def _get_tensor_from_node(self, node_id):
  obj = self._nodes[node_id]
  