text: # Copyright 2018 Google, Inc. All Rights Reserved.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  # ==============================================================================
  
  
  """Evaluation job.
  
  This sits on the side and performs evaluation on a saved model.
  This is a separate process for ease of use and stability of numbers.
  """
  
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  import tensorflow as tf
  
  from learning_unsupervised_learning import utils
  
  
  def construct_evaluation_graph(theta_process_fn=None,
                                 w_learner_fn=None,
                                 dataset_fn=None,
                                 meta_objectives=None,
                                ):
    """Construct the evaluation graph.
    """
    if meta_objectives is None:
      meta_objectives = []
  
    tf.train.create_global_step()
  
    local_device = ""
    remote_device = ""
  
    meta_opt = theta_process_fn(
        remote_device=remote_device, local_device=local_device)
  
    base_model = w_learner_fn(
        remote_device=remote_device, local_device=local_device)
  
    train_dataset = dataset_fn(device=local_device)
  
    # construct variables
    x, outputs = base_model(train_dataset())
    initial_state = base_model.initial_state(meta_opt, max_steps=10)
    next_state = base_model.compute_next_state(outputs, meta_opt, initial_state)
    with utils.state_barrier_context(next_state):
      train_one_step_op = meta_opt.assign_state(base_model, next_state)
  
    meta_objs = []
    for meta_obj_fn in meta_objectives:
      meta_obj = meta_obj_fn(local_device="", remote_device="")
      meta_objs.append(meta_obj)
      J = meta_obj(train_dataset, lambda x: base_model(x)[0])
      tf.summary.scalar(str(meta_obj.__class__.__name__)+"_J", tf.reduce_mean(J))
  
    # TODO(lmetz) this is kinda error prone.
    # We should share the construction of the global variables across train and
    # make sure both sets of savable variables are the same
    checkpoint_vars = meta_opt.remote_variables() + [tf.train.get_global_step()]
    for meta_obj in meta_objs:
      checkpoint_vars.extend(meta_obj.remote_variables())
  
    return checkpoint_vars, train_one_step_op, (base_model, train_dataset)
  