text: # Copyright 2016 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.
  # ==============================================================================
  """Simple examples of the REINFORCE algorithm."""
  
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  import tensorflow as tf
  
  
  distributions = tf.contrib.distributions
  sg = tf.contrib.bayesflow.stochastic_graph
  st = tf.contrib.bayesflow.stochastic_tensor
  
  
  def split_apply_merge(inp, partitions, fns):
    """Split input according to partitions.  Pass results through fns and merge.
  
    Args:
      inp: the input vector
      partitions: tensor of same length as input vector, having values 0, 1
      fns: the two functions.
  
    Returns:
      the vector routed, where routed[i] = fns[partitions[i]](inp[i])
    """
    new_inputs = tf.dynamic_partition(inp, partitions, len(fns))
    new_outputs = [fns[i](x) for i, x in enumerate(new_inputs)]
    new_indices = tf.dynamic_partition(
        tf.range(0, inp.get_shape()[0]), partitions, len(fns))
    return tf.dynamic_stitch(new_indices, new_outputs)
  
  
  def plus_1(inputs):
    return inputs + 1.0
  
  
  def minus_1(inputs):
    return inputs - 1.0
  
  
  def build_split_apply_merge_model():
    """Build the Split-Apply-Merge Model.
  
    Route each value of input [-1, -1, 1, 1] through one of the
    functions, plus_1, minus_1.  The decision for routing is made by
    4 Bernoulli R.V.s whose parameters are determined by a neural network
    applied to the input.  REINFORCE is used to update the NN parameters.
  
    Returns:
      The 3-tuple (route_selection, routing_loss, final_loss), where:
  
        - route_selection is an int 4-vector
        - routing_loss is a float 4-vector
        - final_loss is a float scalar.
    """
    inputs = tf.constant([[-1.0], [-1.0], [1.0], [1.0]])
    targets = tf.constant([[0.0], [0.0], [0.0], [0.0]])
    paths = [plus_1, minus_1]
    weights = tf.get_variable("w", [1, 2])
    bias = tf.get_variable("b", [1, 1])
    logits = tf.matmul(inputs, weights) + bias
  
    # REINFORCE forward step
    route_selection = st.StochasticTensor(
        distributions.Categorical(logits=logits))
  
    # Accessing route_selection as a Tensor below forces a sample of
    # the Categorical distribution based on its logits.
    # This is equivalent to calling route_selection.value().
    #
    # route_selection.value() returns an int32 4-vector with random
    # values in {0, 1}
    # COPY+ROUTE+PASTE
    outputs = split_apply_merge(inputs, route_selection, paths)
  
    # flatten routing_loss to a row vector (from a column vector)
    routing_loss = tf.reshape(tf.square(outputs - targets), shape=[-1])
  
    # Total loss: score function loss + routing loss.
    # The score function loss (through `route_selection.loss(routing_loss)`)
    # returns:
    #  [stop_gradient(routing_loss) *
    #   route_selection.log_pmf(stop_gradient(route_selection.value()))],
    # where log_pmf has gradients going all the way back to weights and bias.
    # In this case, the routing_loss depends on the variables only through
    # "route_selection", which has a stop_gradient on it.  So the
    # gradient of the loss really come through the score function
    surrogate_loss = sg.surrogate_loss([routing_loss])
    final_loss = tf.reduce_sum(surrogate_loss)
  
    return (route_selection, routing_loss, final_loss)
  
  
  class REINFORCESimpleExample(tf.test.TestCase):
  
    def testSplitApplyMerge(self):
      # Repeatability.  SGD has a tendency to jump around, even here.
      tf.set_random_seed(1)
  
      with self.test_session() as sess:
        # Use sampling to train REINFORCE
        with st.value_type(st.SampleValue()):
          (route_selection,
           routing_loss,
           final_loss) = build_split_apply_merge_model()
  
        sgd = tf.train.GradientDescentOptimizer(1.0).minimize(final_loss)
  
        tf.global_variables_initializer().run()
  
        for i in range(10):
          # Run loss and inference step.  This toy problem converges VERY quickly.
          (routing_loss_v, final_loss_v, route_selection_v, _) = sess.run(
              [routing_loss, final_loss, tf.identity(route_selection), sgd])
          print(
              "Iteration %d, routing loss: %s, final_loss: %s, "
              "route selection: %s"
              % (i, routing_loss_v, final_loss_v, route_selection_v))
  
        self.assertAllEqual([0, 0, 1, 1], route_selection_v)
        self.assertAllClose([0.0, 0.0, 0.0, 0.0], routing_loss_v)
        self.assertAllClose(0.0, final_loss_v)
  
  
  if __name__ == "__main__":
    tf.test.main()
  