text: # Copyright 2017 The TensorFlow Authors. 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.
  # ==============================================================================
  """Utilities for Grappler autoparallel optimizer."""
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  import tensorflow as tf
  from tensorflow.core.framework import variable_pb2
  from tensorflow.core.protobuf import rewriter_config_pb2
  
  FLAGS = tf.flags.FLAGS
  
  
  def export_state_tuples(state_tuples, name):
    for state_tuple in state_tuples:
      tf.add_to_collection(name, state_tuple.c)
      tf.add_to_collection(name, state_tuple.h)
  
  
  def import_state_tuples(state_tuples, name, num_replicas):
    restored = []
    for i in range(len(state_tuples) * num_replicas):
      c = tf.get_collection_ref(name)[2 * i + 0]
      h = tf.get_collection_ref(name)[2 * i + 1]
      restored.append(tf.contrib.rnn.LSTMStateTuple(c, h))
    return tuple(restored)
  
  
  def with_prefix(prefix, name):
    """Adds prefix to name."""
    return "/".join((prefix, name))
  
  
  def with_autoparallel_prefix(replica_id, name):
    return with_prefix("AutoParallel-Replica-%d" % replica_id, name)
  
  
  class UpdateCollection(object):
    """Update collection info in MetaGraphDef for AutoParallel optimizer."""
  
    def __init__(self, metagraph, model):
      self._metagraph = metagraph
      self.replicate_states(model.initial_state_name)
      self.replicate_states(model.final_state_name)
      self.update_snapshot_name("variables")
      self.update_snapshot_name("trainable_variables")
  
    def update_snapshot_name(self, var_coll_name):
      var_list = self._metagraph.collection_def[var_coll_name]
      for i, value in enumerate(var_list.bytes_list.value):
        var_def = variable_pb2.VariableDef()
        var_def.ParseFromString(value)
        # Somehow node Model/global_step/read doesn't have any fanout and seems to
        # be only used for snapshot; this is different from all other variables.
        if var_def.snapshot_name != "Model/global_step/read:0":
          var_def.snapshot_name = with_autoparallel_prefix(
              0, var_def.snapshot_name)
        value = var_def.SerializeToString()
        var_list.bytes_list.value[i] = value
  
    def replicate_states(self, state_coll_name):
      state_list = self._metagraph.collection_def[state_coll_name]
      num_states = len(state_list.node_list.value)
      for replica_id in range(1, FLAGS.num_gpus):
        for i in range(num_states):
          state_list.node_list.value.append(state_list.node_list.value[i])
      for replica_id in range(FLAGS.num_gpus):
        for i in range(num_states):
          index = replica_id * num_states + i
          state_list.node_list.value[index] = with_autoparallel_prefix(
              replica_id, state_list.node_list.value[index])
  
  
  def auto_parallel(metagraph, model):
    from tensorflow.python.grappler import tf_optimizer
    rewriter_config = rewriter_config_pb2.RewriterConfig()
    rewriter_config.optimizers.append("autoparallel")
    rewriter_config.auto_parallel.enable = True
    rewriter_config.auto_parallel.num_replicas = FLAGS.num_gpus
    optimized_graph = tf_optimizer.OptimizeGraph(rewriter_config, metagraph)
    metagraph.graph_def.CopyFrom(optimized_graph)
    UpdateCollection(metagraph, model)
  