def testScan_MultiOutputMismatchedInitializer(self):
   elems = np.array([1.0, 2.0, 3.0, 4.0, 5.0, 6.0])
   initializer = np.array(1.0)
   # Multiply a * 1 each time
   with self.assertRaisesRegexp(
       ValueError, "two structures don't have the same nested structure"):
     functional_ops.scan(lambda a, x: (a, -a), elems, initializer)

  def testScan_Simple(self):
    with self.test_session():
      elems = constant_op.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], name="data")
      v = constant_op.constant(2.0, name="v")

      r = functional_ops.scan(lambda a, x: math_ops.mul(a, x), elems)
      self.assertAllEqual([1., 2., 6., 24., 120., 720.], r.eval())

      r = functional_ops.scan(
          lambda a, x: math_ops.mul(a, x), elems, initializer=v)
      self.assertAllEqual([2., 4., 12., 48., 240., 1440.], r.eval())

  def testScan_Simple(self):
    elems = constant_op.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], name="data")
    v = constant_op.constant(2.0, name="v")

    # pylint: disable=unnecessary-lambda
    r = functional_ops.scan(lambda a, x: math_ops.multiply(a, x), elems)
    self.assertAllEqual([1., 2., 6., 24., 120., 720.], self.evaluate(r))

    r = functional_ops.scan(
        lambda a, x: math_ops.multiply(a, x), elems, initializer=v)
    self.assertAllEqual([2., 4., 12., 48., 240., 1440.], self.evaluate(r))

  def testScan_Reverse(self):
    with self.test_session():
      elems = constant_op.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], name="data")
      v = constant_op.constant(2.0, name="v")

      # pylint: disable=unnecessary-lambda
      r = functional_ops.scan(lambda a, x: math_ops.multiply(a, x), elems,
                              reverse=True)
      self.assertAllEqual([720., 720., 360., 120., 30., 6.], self.evaluate(r))
      r = functional_ops.scan(
          lambda a, x: math_ops.multiply(a, x), elems, initializer=v,
          reverse=True)
      self.assertAllEqual([1440., 1440., 720., 240., 60., 12.],
                          self.evaluate(r))

  def testScanVaryingShape(self):
    with self.cached_session() as sess:
      x = array_ops.placeholder(dtype=dtypes.float32, shape=[None, 2])
      x_t = array_ops.transpose(x)
      # scan over dimension 0 (with shape None)
      result = functional_ops.scan(lambda a, x: a + x, x)
      # scanned over transposed dimension 0 (with shape 2)
      result_t = functional_ops.scan(lambda a, x: a + x, x_t, infer_shape=False)
      # ensure gradients can be calculated
      result_grad = gradients_impl.gradients(result, [x])[0]
      result_t_grad = gradients_impl.gradients(result_t, [x_t])[0]

      # smoke test to ensure they all evaluate
      sess.run([result, result_t, result_grad, result_t_grad],
               feed_dict={x: [[1.0, 2.0]]})

 def hiddens(self, input_idxes):
     "Expects input_idxes to be input_idxes of size TIMESTEPS * BATCH_SIZE"
     # embed input encoded sentences
     embedded_timesteps = self.embedding(input_idxes)
     batch_size = tf.shape(input_idxes)[1]
     initial_state = self.rnn_cell.zero_state(batch_size)
     return  functional_ops.scan(self.step_fun, embedded_timesteps, initializer=initial_state)

def power_sums_tensor(array_size, power_matrix, multiplier):
  r"""Computes \sum_{i=0}^{N-1} A^i B (A^i)^T for N=0..(array_size + 1).

  Args:
    array_size: The number of non-trivial sums to pre-compute.
    power_matrix: The "A" matrix above.
    multiplier: The "B" matrix above
  Returns:
    A Tensor with S[N] = \sum_{i=0}^{N-1} A^i B (A^i)^T
      S[0] is the zero matrix
      S[1] is B
      S[2] is A B A^T + B
      ...and so on
  """
  array_size = math_ops.cast(array_size, dtypes.int32)
  power_matrix = ops.convert_to_tensor(power_matrix)
  identity_like_power_matrix = linalg_ops.eye(
      array_ops.shape(power_matrix)[0], dtype=power_matrix.dtype)
  identity_like_power_matrix.set_shape(
      ops.convert_to_tensor(power_matrix).get_shape())
  transition_powers = functional_ops.scan(
      lambda previous_power, _: math_ops.matmul(previous_power, power_matrix),
      math_ops.range(array_size - 1),
      initializer=identity_like_power_matrix)
  summed = math_ops.cumsum(
      array_ops.concat([
          array_ops.expand_dims(multiplier, 0), math_ops.matmul(
              batch_times_matrix(transition_powers, multiplier),
              transition_powers,
              adjoint_b=True)
      ], 0))
  return array_ops.concat(
      [array_ops.expand_dims(array_ops.zeros_like(multiplier), 0), summed], 0)

 def testScan_MultiInputSameTypeOutput(self):
   elems = np.array([1.0, 2.0, 3.0, 4.0, 5.0, 6.0])
   r = functional_ops.scan(lambda a, x: (a[0] + x[0], a[1] + x[1]),
                           (elems, -elems))
   r_value = self.evaluate(r)
   self.assertAllEqual(np.cumsum(elems), r_value[0])
   self.assertAllEqual(np.cumsum(-elems), r_value[1])

 def testScan_MultiInputSingleOutput(self):
   elems = np.array([1.0, 2.0, 3.0, 4.0, 5.0, 6.0])
   initializer = np.array(1.0)
   # Multiply a * 1 each time
   r = functional_ops.scan(lambda a, x: a * (x[0] + x[1]),
                           (elems + 1, -elems), initializer)
   self.assertAllEqual([1.0, 1.0, 1.0, 1.0, 1.0, 1.0], self.evaluate(r))

  def integrate(self, evol_func, y0, time_grid):
    time_delta_grid = time_grid[1:] - time_grid[:-1]

    scan_func = self._make_scan_func(evol_func)

    y_grid = functional_ops.scan(scan_func, (time_grid[:-1], time_delta_grid),
                                 y0)
    return array_ops.concat([[y0], y_grid], axis=0)

  def testScan_SingleInputMultiOutput(self):
    elems = np.array([1.0, 2.0, 3.0, 4.0, 5.0, 6.0])
    initializer = (np.array(1.0), np.array(-1.0))
    r = functional_ops.scan(lambda a, x: (a[0] * x, -a[1] * x), elems,
                            initializer)
    r_value = self.evaluate(r)

    self.assertAllEqual([1.0, 2.0, 6.0, 24.0, 120.0, 720.0], r_value[0])
    self.assertAllEqual([1.0, -2.0, 6.0, -24.0, 120.0, -720.0], r_value[1])

  def testScanShape(self):
    x = constant_op.constant([[1, 2, 3], [4, 5, 6]])

    def fn(_, current_input):
      return current_input

    initializer = constant_op.constant([0, 0, 0])
    y = functional_ops.scan(fn, x, initializer=initializer)
    self.assertAllEqual(y.get_shape(), self.evaluate(y).shape)

  def testScanUnknownShape(self):
    x = array_ops.placeholder(dtypes.float32)
    initializer = array_ops.placeholder(dtypes.float32)

    def fn(_, current_input):
      return current_input

    y = functional_ops.scan(fn, x, initializer=initializer)
    self.assertIs(None, y.get_shape().dims)

  def testScan_Grad(self):
    with self.test_session():
      elems = constant_op.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], name="data")
      v = constant_op.constant(2.0, name="v")

      r = functional_ops.scan(
          lambda a, x: math_ops.mul(a, x), elems, initializer=v)
      r = gradients_impl.gradients(r, v)[0]
      self.assertAllEqual(873.0, r.eval())

    def _compute_predictions(self, init = None):
        """ Compute vanilla-RNN states and predictions. """

        with tf.variable_scope('states'):
            with tf.variable_scope("HMM"):
                with tf.variable_scope("transition"):
                    skip_prob = tf.get_variable("skip", shape=[1], initializer=tf.constant_initializer(1e-1))
                    #skip_prob = tf.Variable( np.array(1e-1, dtype=np.float32), name="skip") # .astype(np.float32)
                    self.W_trans = (1-skip_prob) * get_transition_matrix().astype(np.float32)  + skip_prob* np.eye(self.hidden_layer_size).astype(np.float32)
                    #self.W_trans = tf.Variable( transition_with_skips,
                    #                       name='W_trans', trainable=True)
                    print("W_trans", self.W_trans.get_shape())

                with tf.variable_scope("emission"):
                    "W_emit: [self.input_size, self.hidden_layer_size]"
                    if self.emission_init is None:
                        self.W_emit = tf.get_variable("W_emit", shape = [self.hidden_layer_size, self.input_size],
                                                  initializer = tf.random_normal_initializer(0.0, 1e-6))
                    else:
                        if not (self.emission_init.shape == (self.hidden_layer_size, self.input_size)):
                            print("self.emission_init.shape", self.emission_init.shape)
                            print("(self.hidden_layer_size, self.input_size)", (self.hidden_layer_size, self.input_size))
                            raise ValueError("wrong dimensions of  `self.emission_init`")
                        self.W_emit = tf.Variable(self.emission_init.astype(np.float32), name = "W_emit", trainable = False)
                    self.W_emit_summary = tf.image_summary("W_emit", tf.reshape(self.W_emit, [1,self.hidden_layer_size, self.input_size,1]))
                    "idea: impose kernel similarity:  maximize(W K W)"
                    "[ self.hidden_layer_size, self.nt_in_pore ]"

                    emission_in_pore_space = tf.matmul( self.map_hex_to_pore, self.W_emit)
                    self.emission_similarity = tf.reduce_sum( tf.diag_part( tf.matmul( tf.transpose(emission_in_pore_space),(emission_in_pore_space)) ),
                            name="emission_w_similarity")
            if init is None:
                initial_state = tf.ones([self.hidden_layer_size],
                                     name='initial_state')
                initial_state = initial_state/ self.hidden_layer_size
            else:
                initial_state = init
            #states = self._rnn_step_fw(initial_state[:,0], self.inputs[0,:])
            states = functional_ops.scan(self._rnn_step_fw, tf.identity(self.inputs),
                                         initializer=initial_state, name='states')

            states_fw_summary = tf.histogram_summary("states_fw", states)
            #states = states_fw
            #print("states:", states.get_shape())

        with tf.variable_scope('predictions'):
            # set some explicit initializer, orthogonal inialization
            "for now, keep identity mapping from hidden states to labels"
            "assume probability interpretation of values: should sum to one"
            W_pred = tf.Variable(np.eye(self.target_size, dtype = np.float32), name="W_pred", trainable=False)
            predictions = tf.matmul(states, W_pred, name='predictions')
            #predictions = states
            predictions_summary = tf.histogram_summary("predictions", predictions)
            #predictions = tf.nn.softmax(tf.matmul(states, W_pred), name='predictions'))
            # do predictions sum to one?

        return states, predictions

  def test_jacobian_scan_shape(self):
    # Shape x: [3, 4]
    x = random_ops.random_uniform([3, 4])
    elems = random_ops.random_uniform([6])
    # Shape y: [6, 3, 4]
    y = functional_ops.scan(lambda a, e: a + e, elems, initializer=x)
    jacobian = gradients.jacobian(y, x)

    expected_shape = [6, 3, 4, 3, 4]
    self.assertAllEqual(expected_shape, jacobian.shape.as_list())

  def testScan_Control(self):
    with self.cached_session() as sess:
      s = array_ops.placeholder(dtypes.float32, shape=[None])
      b = array_ops.placeholder(dtypes.bool)

      with ops.control_dependencies([b]):
        c = functional_ops.scan(lambda a, x: x * a, s)
      self.assertAllClose(
          np.array([1.0, 3.0, 9.0]), sess.run(c, {s: [1, 3, 3],
                                                  b: True}))

  def testScan_Scoped(self):
    with self.cached_session() as sess:
      with variable_scope.variable_scope("root") as varscope:
        elems = constant_op.constant([1, 2, 3, 4, 5, 6], name="data")

        r = functional_ops.scan(simple_scoped_fn, elems)
        # Check that we have the one variable we asked for here.
        self.assertEqual(len(variables.trainable_variables()), 1)
        self.assertEqual(variables.trainable_variables()[0].name,
                         "root/body/two:0")
        sess.run([variables.global_variables_initializer()])
        results = np.array([1, 6, 18, 44, 98, 208])
        self.assertAllEqual(results, self.evaluate(r))

        # Now let's reuse our single variable.
        varscope.reuse_variables()
        r = functional_ops.scan(simple_scoped_fn, elems, initializer=2)
        self.assertEqual(len(variables.trainable_variables()), 1)
        results = np.array([6, 16, 38, 84, 178, 368])
        self.assertAllEqual(results, self.evaluate(r))

  def testScan_Grad(self):
    with self.cached_session():
      elems = constant_op.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], name="data")
      v = constant_op.constant(2.0, name="v")

      # pylint: disable=unnecessary-lambda
      r = functional_ops.scan(
          lambda a, x: math_ops.multiply(a, x), elems, initializer=v)
      # pylint: enable=unnecessary-lambda
      r = gradients_impl.gradients(r, v)[0]
      self.assertAllEqual(873.0, self.evaluate(r))

 def testScanGradientWithPartStopGradient(self):
   a = variables.Variable(0.0, name="a")
   b = variables.Variable(0.0, name="b")
   elems = array_ops.zeros(5)
   l0, l1 = functional_ops.scan(
       lambda elem_, input_: (a, b), elems, initializer=(0., 0.))
   loss = l0 + array_ops.stop_gradient(l1)
   grad = gradients_impl.gradients(ys=[loss], xs=[a, b])
   with self.test_session(use_gpu=True) as sess:
     self.evaluate(variables.global_variables_initializer())
     self.evaluate(grad)