Python tensorflow.assign_add() Examples

def _graph_fn_get_should_sync(self):
        if get_backend() == "tf":
            inc_op = tf.assign_add(self.steps_since_last_sync, 1)
            should_sync = inc_op >= self.q_sync_spec.sync_interval

            def reset_op():
                op = tf.assign(self.steps_since_last_sync, 0)
                with tf.control_dependencies([op]):
                    return tf.no_op()

            sync_op = tf.cond(
                pred=inc_op >= self.q_sync_spec.sync_interval,
                true_fn=reset_op,
                false_fn=tf.no_op
            )
            with tf.control_dependencies([sync_op]):
                return tf.identity(should_sync)
        else:
            raise NotImplementedError("TODO") 

def apply_stats_eigen(self, eigen_list):
        """
        apply the update using the eigen values of the stats

        :param eigen_list: ([TensorFlow Tensor]) The list of eigen values of the stats
        :return: ([TensorFlow Tensor]) update operations
        """
        update_ops = []
        if self.verbose > 1:
            print(('updating %d eigenvalue/vectors' % len(eigen_list)))
        for _, (tensor, mark) in enumerate(zip(eigen_list, self.eigen_update_list)):
            stats_eigen_var = self.eigen_reverse_lookup[mark]
            update_ops.append(
                tf.assign(stats_eigen_var, tensor, use_locking=True))

        with tf.control_dependencies(update_ops):
            factor_step_op = tf.assign_add(self.factor_step, 1)
            update_ops.append(factor_step_op)
            if KFAC_DEBUG:
                update_ops.append(tf.Print(tf.constant(
                    0.), [tf.convert_to_tensor('updated kfac factors')]))
        return update_ops 

def _apply_cond(self, apply_fn, grad, var, *args, **kwargs):
    """Apply conditionally if counter is zero."""
    grad_acc = self.get_slot(var, "grad_acc")

    def apply_adam(grad_acc, apply_fn, grad, var, *args, **kwargs):
      total_grad = (grad_acc + grad) / tf.cast(self._n_t, grad.dtype)
      adam_op = apply_fn(total_grad, var, *args, **kwargs)
      with tf.control_dependencies([adam_op]):
        grad_acc_to_zero_op = grad_acc.assign(tf.zeros_like(grad_acc),
                                              use_locking=self._use_locking)
      return tf.group(adam_op, grad_acc_to_zero_op)

    def accumulate_gradient(grad_acc, grad):
      assign_op = tf.assign_add(grad_acc, grad, use_locking=self._use_locking)
      return tf.group(assign_op)  # Strip return value

    return tf.cond(
        tf.equal(self._get_iter_variable(), 0),
        lambda: apply_adam(grad_acc, apply_fn, grad, var, *args, **kwargs),
        lambda: accumulate_gradient(grad_acc, grad)) 

def call_fake_controller(push_values, pop_values, read_values, output_values):
  """Mock a RNN controller from a set of expected outputs.

  Args:
    push_values: Expected controller push values.
    pop_values: Expected controller pop values.
    read_values: Expected controller read values.
    output_values: Expected controller output values.

  Returns:
    A callable which behaves like the call method of an NeuralStackCell.
  """
  def call(cell, inputs, state, batch_size):
    del inputs
    del batch_size
    next_step = tf.assign_add(cell.current_step, tf.constant(1))
    return (
        tf.slice(tf.constant(push_values), [next_step, 0], [1, -1]),
        tf.slice(tf.constant(pop_values), [next_step, 0], [1, -1]),
        tf.slice(tf.constant(read_values), [next_step, 0, 0], [1, -1, -1]),
        tf.slice(tf.constant(output_values), [next_step, 0, 0], [1, -1, -1]),
        state
    )
  return call 

def test_ref_assign_add(self):
        # Currently ngraph and tf have different assign semantics
        # eval(ng.assign(a, 1)) resturns None, but eval(tf.assign(a, 1)) returns
        # a which is 1.
        # TODO: fix this test after assign op / user_deps are fixed in ngraph
        # TODO: double assignments fails

        # tf placeholder
        a = tf.Variable(tf.constant(np.random.randn(2, 3), name="a"))
        b = tf.Variable(tf.constant(np.random.randn(2, 3), name="b"))
        init_op = tf.global_variables_initializer()
        a_update = tf.assign_add(a, b)

        # test
        tf_result = self.tf_run(a_update, tf_init_op=init_op)
        ng_result = self.ng_run(a)
        ng.testing.assert_allclose(tf_result, ng_result) 

def get_action(self):
        deterministic_action = self.mean

        if self.ornstein_uhlenbeck_state:
            random_norm = tf.random_normal(
                shape=self.mean.get_shape().as_list()[1:],
                mean=0.0,
                stddev=1.0)
            ou_noise = tf.assign_add(
                self.ornstein_uhlenbeck_state,
                self._theta * (-self.ornstein_uhlenbeck_state) +
                random_norm * self._sigma)
            sample_action = self.mean + ou_noise * self._noise_scale

            output_action = tf.cond(self.deterministic_ph,
                                    lambda: deterministic_action,
                                    lambda: sample_action)
        else:
            output_action = deterministic_action

        return output_action 

def _load_sampler(self):
    if not self.structure_path:
      print("Loading the default file")
      with open('../../../../bayesian_optimization/sample_structures6.txt', 'r') as fp:
        lines = fp.readlines()
    else:
      with open(self.structure_path, 'r') as fp:
        lines = fp.readlines()
    lines = [eval(line.strip()) for line in lines]
    for line in lines:
      row = []
      for ele in line:
        row += ele
      self.structures.append(row)
    # self.sample_arcs = self.structures[0]
    self.idx_arc = tf.get_variable("idx_arc", initializer=tf.constant(0), dtype=tf.int32)
    self.reset_op = self.idx_arc.assign(0)
    self.inc_op = self.idx_arc.assign(self.idx_arc + 1)
    # self.idx_arc  = tf.assign_add(self.idx_arc, tf.constant(1, dtype=tf.int32))
    self.sample_arc2 = tf.reshape(tf.gather(self.structures, self.idx_arc), [-1])
    # self.sample_arc3 = tf.get_variable("sample_arc3", shape=[len(self.structures[0])])
    self.sample_arc3 = tf.placeholder(tf.float32, shape=[len(self.structures[0])]) 

def build_act(actor_num, env, q_func, num_actions, eps, scope="deepq", data_format=None, reuse=None, replay_queue=None, prioritized_replay_eps=None, gamma=None, replay_queue_capacity=None, multi_step_n=1, framestack=4, num_actor_steps=None):
    # Build the actor using \epsilon-greedy exploration
    # Environments are wrapped with a data gathering class
    with tf.variable_scope('deepq', reuse=reuse):
        with tf.device('/gpu:0'):
            act_obs = env.observation()
            q_values = q_func(act_obs, num_actions, scope="read_q_func", data_format=data_format)
        deterministic_actions = tf.argmax(q_values, axis=1, output_type=tf.int32)
        batch_size = deterministic_actions.get_shape()[0]
        with tf.device('/cpu:0'):
            random_actions = tf.random_uniform(tf.stack([batch_size]), minval=0, maxval=num_actions, dtype=tf.int32)
            chose_random = tf.random_uniform(tf.stack([batch_size]), minval=0, maxval=1, dtype=tf.float32) < eps
            output_actions = tf.where(chose_random, random_actions, deterministic_actions)

            q_t_selected = tf.reduce_sum(q_values * tf.one_hot(output_actions, num_actions), 1)

        with tf.control_dependencies([tf.assign_add(num_actor_steps, batch_size, use_locking=True)]):
            return env.step(output_actions, q_values=q_values, q_t_selected=q_t_selected) 

def testBasicDomainSeparationStartPoint(self):
    with self.test_session() as sess:
      # Test for when global_step < domain_separation_startpoint
      step = tf.contrib.slim.get_or_create_global_step()
      sess.run(tf.global_variables_initializer())  # global_step = 0
      params = {'domain_separation_startpoint': 2}
      weight = dsn.dsn_loss_coefficient(params)
      weight_np = sess.run(weight)
      self.assertAlmostEqual(weight_np, 1e-10)

      step_op = tf.assign_add(step, 1)
      step_np = sess.run(step_op)  # global_step = 1
      weight = dsn.dsn_loss_coefficient(params)
      weight_np = sess.run(weight)
      self.assertAlmostEqual(weight_np, 1e-10)

      # Test for when global_step >= domain_separation_startpoint
      step_np = sess.run(step_op)  # global_step = 2
      tf.logging.info(step_np)
      weight = dsn.dsn_loss_coefficient(params)
      weight_np = sess.run(weight)
      self.assertAlmostEqual(weight_np, 1.0) 

def _create_autosummary_var(name, value_expr):
    assert not _autosummary_finalized
    v = tf.cast(value_expr, tf.float32)
    if v.shape.ndims is 0:
        v = [v, np.float32(1.0)]
    elif v.shape.ndims is 1:
        v = [tf.reduce_sum(v), tf.cast(tf.shape(v)[0], tf.float32)]
    else:
        v = [tf.reduce_sum(v), tf.reduce_prod(tf.cast(tf.shape(v), tf.float32))]
    v = tf.cond(tf.is_finite(v[0]), lambda: tf.stack(v), lambda: tf.zeros(2))
    with tf.control_dependencies(None):
        var = tf.Variable(tf.zeros(2)) # [numerator, denominator]
    update_op = tf.cond(tf.is_variable_initialized(var), lambda: tf.assign_add(var, v), lambda: tf.assign(var, v))
    if name in _autosummary_vars:
        _autosummary_vars[name].append(var)
    else:
        _autosummary_vars[name] = [var]
    return update_op

#----------------------------------------------------------------------------
# Call filewriter.add_summary() with all summaries in the default graph,
# automatically finalizing and merging them on the first call. 

def apply_gradients(self, grads):
        coldOptim = tf.train.MomentumOptimizer(
            self._cold_lr, self._momentum)

        def coldSGDstart():
            sgd_grads, sgd_var = zip(*grads)

            if self.max_grad_norm != None:
                sgd_grads, sgd_grad_norm = tf.clip_by_global_norm(sgd_grads,self.max_grad_norm)

            sgd_grads = list(zip(sgd_grads,sgd_var))

            sgd_step_op = tf.assign_add(self.sgd_step, 1)
            coldOptim_op = coldOptim.apply_gradients(sgd_grads)
            if KFAC_DEBUG:
                with tf.control_dependencies([sgd_step_op, coldOptim_op]):
                    sgd_step_op = tf.Print(
                        sgd_step_op, [self.sgd_step, tf.convert_to_tensor('doing cold sgd step')])
            return tf.group(*[sgd_step_op, coldOptim_op])

        kfacOptim_op, qr = self.apply_gradients_kfac(grads)

        def warmKFACstart():
            return kfacOptim_op

        return tf.cond(tf.greater(self.sgd_step, self._cold_iter), warmKFACstart, coldSGDstart), qr 

def applyStatsEigen(self, eigen_list):
        updateOps = []
        print(('updating %d eigenvalue/vectors' % len(eigen_list)))
        for i, (tensor, mark) in enumerate(zip(eigen_list, self.eigen_update_list)):
            stats_eigen_var = self.eigen_reverse_lookup[mark]
            updateOps.append(
                tf.assign(stats_eigen_var, tensor, use_locking=True))

        with tf.control_dependencies(updateOps):
            factor_step_op = tf.assign_add(self.factor_step, 1)
            updateOps.append(factor_step_op)
            if KFAC_DEBUG:
                updateOps.append(tf.Print(tf.constant(
                    0.), [tf.convert_to_tensor('updated kfac factors')]))
        return updateOps 

def _apply_stats(self, stats_updates, accumulate=False, accumulate_coeff=0.):
        update_ops = []
        # obtain the stats var list
        for stats_var in stats_updates:
            stats_new = stats_updates[stats_var]
            if accumulate:
                # simple superbatch averaging
                update_op = tf.assign_add(
                    stats_var, accumulate_coeff * stats_new, use_locking=True)
            else:
                # exponential running averaging
                update_op = tf.assign(
                    stats_var, stats_var * self._stats_decay, use_locking=True)
                update_op = tf.assign_add(
                    update_op, (1. - self._stats_decay) * stats_new, use_locking=True)
            update_ops.append(update_op)

        with tf.control_dependencies(update_ops):
            stats_step_op = tf.assign_add(self.stats_step, 1)

        if KFAC_DEBUG:
            stats_step_op = (tf.Print(stats_step_op,
                                      [tf.convert_to_tensor('step:'),
                                       self.global_step,
                                       tf.convert_to_tensor('fac step:'),
                                       self.factor_step,
                                       tf.convert_to_tensor('sgd step:'),
                                       self.sgd_step,
                                       tf.convert_to_tensor('Accum:'),
                                       tf.convert_to_tensor(accumulate),
                                       tf.convert_to_tensor('Accum coeff:'),
                                       tf.convert_to_tensor(accumulate_coeff),
                                       tf.convert_to_tensor('stat step:'),
                                       self.stats_step, update_ops[0], update_ops[1]]))
        return [stats_step_op, ] 

def applyStatsEigen(self, eigen_list):
        updateOps = []
        print(('updating %d eigenvalue/vectors' % len(eigen_list)))
        for i, (tensor, mark) in enumerate(zip(eigen_list, self.eigen_update_list)):
            stats_eigen_var = self.eigen_reverse_lookup[mark]
            updateOps.append(
                tf.assign(stats_eigen_var, tensor, use_locking=True))

        with tf.control_dependencies(updateOps):
            factor_step_op = tf.assign_add(self.factor_step, 1)
            updateOps.append(factor_step_op)
            if KFAC_DEBUG:
                updateOps.append(tf.Print(tf.constant(
                    0.), [tf.convert_to_tensor('updated kfac factors')]))
        return updateOps 

def apply_gradients(self, grads):
        """
        apply the gradient

        :param grads: ([TensorFlow Tensor]) the gradient
        :return: (function, QueueRunner) train operation, queue operation runner
        """
        cold_optim = tf.train.MomentumOptimizer(self._cold_lr, self._momentum)

        def _cold_sgd_start():
            sgd_grads, sgd_var = zip(*grads)

            if self.max_grad_norm is not None:
                sgd_grads, _ = tf.clip_by_global_norm(sgd_grads, self.max_grad_norm)

            sgd_grads = list(zip(sgd_grads, sgd_var))

            sgd_step_op = tf.assign_add(self.sgd_step, 1)
            cold_optim_op = cold_optim.apply_gradients(sgd_grads)
            if KFAC_DEBUG:
                with tf.control_dependencies([sgd_step_op, cold_optim_op]):
                    sgd_step_op = tf.Print(
                        sgd_step_op, [self.sgd_step, tf.convert_to_tensor('doing cold sgd step')])
            return tf.group(*[sgd_step_op, cold_optim_op])

        # remove unused variables
        grads = [(grad, var) for (grad, var) in grads if grad is not None]

        kfac_optim_op, queue_runner = self.apply_gradients_kfac(grads)

        def _warm_kfac_start():
            return kfac_optim_op

        return tf.cond(tf.greater(self.sgd_step, self._cold_iter), _warm_kfac_start, _cold_sgd_start), queue_runner 

def __init__(self, epsilon=1e-2, shape=()):

        self._sum = tf.get_variable(
            dtype=tf.float64,
            shape=shape,
            initializer=tf.constant_initializer(0.0),
            name="runningsum", trainable=False)
        self._sumsq = tf.get_variable(
            dtype=tf.float64,
            shape=shape,
            initializer=tf.constant_initializer(epsilon),
            name="runningsumsq", trainable=False)
        self._count = tf.get_variable(
            dtype=tf.float64,
            shape=(),
            initializer=tf.constant_initializer(epsilon),
            name="count", trainable=False)
        self.shape = shape

        self.mean = tf.to_float(self._sum / self._count)
        self.std = tf.sqrt( tf.maximum( tf.to_float(self._sumsq / self._count) - tf.square(self.mean) , 1e-2 ))

        newsum = tf.placeholder(shape=self.shape, dtype=tf.float64, name='sum')
        newsumsq = tf.placeholder(shape=self.shape, dtype=tf.float64, name='var')
        newcount = tf.placeholder(shape=[], dtype=tf.float64, name='count')
        self.incfiltparams = U.function([newsum, newsumsq, newcount], [],
            updates=[tf.assign_add(self._sum, newsum),
                     tf.assign_add(self._sumsq, newsumsq),
                     tf.assign_add(self._count, newcount)]) 

def _apply_stats(self, statsUpdates, accumulate=False, accumulateCoeff=0.):
        updateOps = []
        # obtain the stats var list
        for stats_var in statsUpdates:
            stats_new = statsUpdates[stats_var]
            if accumulate:
                # simple superbatch averaging
                update_op = tf.assign_add(
                    stats_var, accumulateCoeff * stats_new, use_locking=True)
            else:
                # exponential running averaging
                update_op = tf.assign(
                    stats_var, stats_var * self._stats_decay, use_locking=True)
                update_op = tf.assign_add(
                    update_op, (1. - self._stats_decay) * stats_new, use_locking=True)
            updateOps.append(update_op)

        with tf.control_dependencies(updateOps):
            stats_step_op = tf.assign_add(self.stats_step, 1)

        if KFAC_DEBUG:
            stats_step_op = (tf.Print(stats_step_op,
                                      [tf.convert_to_tensor('step:'),
                                       self.global_step,
                                       tf.convert_to_tensor('fac step:'),
                                       self.factor_step,
                                       tf.convert_to_tensor('sgd step:'),
                                       self.sgd_step,
                                       tf.convert_to_tensor('Accum:'),
                                       tf.convert_to_tensor(accumulate),
                                       tf.convert_to_tensor('Accum coeff:'),
                                       tf.convert_to_tensor(accumulateCoeff),
                                       tf.convert_to_tensor('stat step:'),
                                       self.stats_step, updateOps[0], updateOps[1]]))
        return [stats_step_op, ] 

def __init__(self, epsilon=1e-2, shape=()):

        self._sum = tf.get_variable(
            dtype=tf.float64,
            shape=shape,
            initializer=tf.constant_initializer(0.0),
            name="runningsum", trainable=False)
        self._sumsq = tf.get_variable(
            dtype=tf.float64,
            shape=shape,
            initializer=tf.constant_initializer(epsilon),
            name="runningsumsq", trainable=False)
        self._count = tf.get_variable(
            dtype=tf.float64,
            shape=(),
            initializer=tf.constant_initializer(epsilon),
            name="count", trainable=False)
        self.shape = shape

        self.mean = tf.to_float(self._sum / self._count)
        self.std = tf.sqrt( tf.maximum( tf.to_float(self._sumsq / self._count) - tf.square(self.mean) , 1e-2 ))

        newsum = tf.placeholder(shape=self.shape, dtype=tf.float64, name='sum')
        newsumsq = tf.placeholder(shape=self.shape, dtype=tf.float64, name='var')
        newcount = tf.placeholder(shape=[], dtype=tf.float64, name='count')
        self.incfiltparams = U.function([newsum, newsumsq, newcount], [],
            updates=[tf.assign_add(self._sum, newsum),
                     tf.assign_add(self._sumsq, newsumsq),
                     tf.assign_add(self._count, newcount)]) 

def _apply_stats(self, statsUpdates, accumulate=False, accumulateCoeff=0.):
        updateOps = []
        # obtain the stats var list
        for stats_var in statsUpdates:
            stats_new = statsUpdates[stats_var]
            if accumulate:
                # simple superbatch averaging
                update_op = tf.assign_add(
                    stats_var, accumulateCoeff * stats_new, use_locking=True)
            else:
                # exponential running averaging
                update_op = tf.assign(
                    stats_var, stats_var * self._stats_decay, use_locking=True)
                update_op = tf.assign_add(
                    update_op, (1. - self._stats_decay) * stats_new, use_locking=True)
            updateOps.append(update_op)

        with tf.control_dependencies(updateOps):
            stats_step_op = tf.assign_add(self.stats_step, 1)

        if KFAC_DEBUG:
            stats_step_op = (tf.Print(stats_step_op,
                                      [tf.convert_to_tensor('step:'),
                                       self.global_step,
                                       tf.convert_to_tensor('fac step:'),
                                       self.factor_step,
                                       tf.convert_to_tensor('sgd step:'),
                                       self.sgd_step,
                                       tf.convert_to_tensor('Accum:'),
                                       tf.convert_to_tensor(accumulate),
                                       tf.convert_to_tensor('Accum coeff:'),
                                       tf.convert_to_tensor(accumulateCoeff),
                                       tf.convert_to_tensor('stat step:'),
                                       self.stats_step, updateOps[0], updateOps[1]]))
        return [stats_step_op, ] 

def _apply_stats(self, statsUpdates, accumulate=False, accumulateCoeff=0.):
        updateOps = []
        # obtain the stats var list
        for stats_var in statsUpdates:
            stats_new = statsUpdates[stats_var]
            if accumulate:
                # simple superbatch averaging
                update_op = tf.assign_add(
                    stats_var, accumulateCoeff * stats_new, use_locking=True)
            else:
                # exponential running averaging
                update_op = tf.assign(
                    stats_var, stats_var * self._stats_decay, use_locking=True)
                update_op = tf.assign_add(
                    update_op, (1. - self._stats_decay) * stats_new, use_locking=True)
            updateOps.append(update_op)

        with tf.control_dependencies(updateOps):
            stats_step_op = tf.assign_add(self.stats_step, 1)

        if KFAC_DEBUG:
            stats_step_op = (tf.Print(stats_step_op,
                                      [tf.convert_to_tensor('step:'),
                                       self.global_step,
                                       tf.convert_to_tensor('fac step:'),
                                       self.factor_step,
                                       tf.convert_to_tensor('sgd step:'),
                                       self.sgd_step,
                                       tf.convert_to_tensor('Accum:'),
                                       tf.convert_to_tensor(accumulate),
                                       tf.convert_to_tensor('Accum coeff:'),
                                       tf.convert_to_tensor(accumulateCoeff),
                                       tf.convert_to_tensor('stat step:'),
                                       self.stats_step, updateOps[0], updateOps[1]]))
        return [stats_step_op, ] 

def testStop(self):
    global_step = tf.train.create_global_step()
    tf.summary.scalar("global_step", global_step)
    incr_global_step = tf.assign_add(global_step, 1)

    ckpt_dir = self.ckpt_dir("stop")
    dummy = DummyHook(ckpt_dir, every_n_steps=10)
    with self.sess(dummy, ckpt_dir) as sess:
      for _ in range(20):
        sess.run(incr_global_step)

      # Summary files should now have 2 global step values in them
      self.flush()

      # Run for 10 more so that the hook gets triggered again
      for _ in range(10):
        sess.run(incr_global_step)

      # Check that the metrics have actually been collected.
      self.assertTrue("" in dummy.test_metrics)
      metrics = dummy.test_metrics[""]
      self.assertTrue("global_step_1" in metrics)
      steps, vals = metrics["global_step_1"]
      self.assertTrue(len(steps) == len(vals))
      self.assertTrue(len(steps) >= 2)

      # Run for 10 more so that the hook triggers stoppage
      for _ in range(10):
        sess.run(incr_global_step)

      with self.assertRaisesRegexp(RuntimeError, "after should_stop requested"):
        sess.run(incr_global_step) 

def _graph_fn_get_noise(self):
        drift = self.theta * (self.mu - self.ou_state)
        if get_backend() == "tf":
            diffusion = self.sigma * tf.random_normal(
                shape=self.action_space.shape, dtype=convert_dtype(self.action_space.dtype)
            )
            delta = drift + diffusion
            return tf.assign_add(ref=self.ou_state, value=delta) 

def _graph_fn_update_global_timestep(self, increment):
        if get_backend() == "tf":
            add_op = tf.assign_add(self.agent.graph_executor.global_timestep, increment)
            return add_op
        elif get_backend == "pytorch":
            self.agent.graph_executor.global_timestep += increment
            return self.agent.graph_executor.global_timestep 

def _graph_fn_training_step(self, other_step_op=None):
        if self.agent is not None:
            add_op = tf.assign_add(self.agent.graph_executor.global_training_timestep, 1)
            op_list = [add_op] + [other_step_op] if other_step_op is not None else []
            with tf.control_dependencies(op_list):
                return tf.no_op() if other_step_op is None else other_step_op
        else:
            return tf.no_op() if other_step_op is None else other_step_op 

def _create_var(name: str, value_expr: TfExpression) -> TfExpression:
    """Internal helper for creating autosummary accumulators."""
    assert not _finalized
    name_id = name.replace("/", "_")
    v = tf.cast(value_expr, _dtype)

    if v.shape.is_fully_defined():
        size = np.prod(tfutil.shape_to_list(v.shape))
        size_expr = tf.constant(size, dtype=_dtype)
    else:
        size = None
        size_expr = tf.reduce_prod(tf.cast(tf.shape(v), _dtype))

    if size == 1:
        if v.shape.ndims != 0:
            v = tf.reshape(v, [])
        v = [size_expr, v, tf.square(v)]
    else:
        v = [size_expr, tf.reduce_sum(v), tf.reduce_sum(tf.square(v))]
    v = tf.cond(tf.is_finite(v[1]), lambda: tf.stack(v), lambda: tf.zeros(3, dtype=_dtype))

    with tfutil.absolute_name_scope("Autosummary/" + name_id), tf.control_dependencies(None):
        var = tf.Variable(tf.zeros(3, dtype=_dtype), trainable=False)  # [sum(1), sum(x), sum(x**2)]
    update_op = tf.cond(tf.is_variable_initialized(var), lambda: tf.assign_add(var, v), lambda: tf.assign(var, v))

    if name in _vars:
        _vars[name].append(var)
    else:
        _vars[name] = [var]
    return update_op 

def __init__(self, epsilon=1e-2, shape=()):

        self._sum = tf.get_variable(
            dtype=tf.float64,
            shape=shape,
            initializer=tf.constant_initializer(0.0),
            name="runningsum", trainable=False)
        self._sumsq = tf.get_variable(
            dtype=tf.float64,
            shape=shape,
            initializer=tf.constant_initializer(epsilon),
            name="runningsumsq", trainable=False)
        self._count = tf.get_variable(
            dtype=tf.float64,
            shape=(),
            initializer=tf.constant_initializer(epsilon),
            name="count", trainable=False)
        self.shape = shape

        self.mean = tf.to_float(self._sum / self._count)
        self.std = tf.sqrt( tf.maximum( tf.to_float(self._sumsq / self._count) - tf.square(self.mean) , 1e-2 ))

        newsum = tf.placeholder(shape=self.shape, dtype=tf.float64, name='sum')
        newsumsq = tf.placeholder(shape=self.shape, dtype=tf.float64, name='var')
        newcount = tf.placeholder(shape=[], dtype=tf.float64, name='count')
        self.incfiltparams = U.function([newsum, newsumsq, newcount], [],
            updates=[tf.assign_add(self._sum, newsum),
                     tf.assign_add(self._sumsq, newsumsq),
                     tf.assign_add(self._count, newcount)]) 

def apply_gradients(self, grads):
        coldOptim = tf.train.MomentumOptimizer(
            self._cold_lr, self._momentum)

        def coldSGDstart():
            sgd_grads, sgd_var = zip(*grads)

            if self.max_grad_norm != None:
                sgd_grads, sgd_grad_norm = tf.clip_by_global_norm(sgd_grads,self.max_grad_norm)

            sgd_grads = list(zip(sgd_grads,sgd_var))

            sgd_step_op = tf.assign_add(self.sgd_step, 1)
            coldOptim_op = coldOptim.apply_gradients(sgd_grads)
            if KFAC_DEBUG:
                with tf.control_dependencies([sgd_step_op, coldOptim_op]):
                    sgd_step_op = tf.Print(
                        sgd_step_op, [self.sgd_step, tf.convert_to_tensor('doing cold sgd step')])
            return tf.group(*[sgd_step_op, coldOptim_op])

        kfacOptim_op, qr = self.apply_gradients_kfac(grads)

        def warmKFACstart():
            return kfacOptim_op

        return tf.cond(tf.greater(self.sgd_step, self._cold_iter), warmKFACstart, coldSGDstart), qr 

def applyStatsEigen(self, eigen_list):
        updateOps = []
        print(('updating %d eigenvalue/vectors' % len(eigen_list)))
        for i, (tensor, mark) in enumerate(zip(eigen_list, self.eigen_update_list)):
            stats_eigen_var = self.eigen_reverse_lookup[mark]
            updateOps.append(
                tf.assign(stats_eigen_var, tensor, use_locking=True))

        with tf.control_dependencies(updateOps):
            factor_step_op = tf.assign_add(self.factor_step, 1)
            updateOps.append(factor_step_op)
            if KFAC_DEBUG:
                updateOps.append(tf.Print(tf.constant(
                    0.), [tf.convert_to_tensor('updated kfac factors')]))
        return updateOps 

def _apply_stats(self, statsUpdates, accumulate=False, accumulateCoeff=0.):
        updateOps = []
        # obtain the stats var list
        for stats_var in statsUpdates:
            stats_new = statsUpdates[stats_var]
            if accumulate:
                # simple superbatch averaging
                update_op = tf.assign_add(
                    stats_var, accumulateCoeff * stats_new, use_locking=True)
            else:
                # exponential running averaging
                update_op = tf.assign(
                    stats_var, stats_var * self._stats_decay, use_locking=True)
                update_op = tf.assign_add(
                    update_op, (1. - self._stats_decay) * stats_new, use_locking=True)
            updateOps.append(update_op)

        with tf.control_dependencies(updateOps):
            stats_step_op = tf.assign_add(self.stats_step, 1)

        if KFAC_DEBUG:
            stats_step_op = (tf.Print(stats_step_op,
                                      [tf.convert_to_tensor('step:'),
                                       self.global_step,
                                       tf.convert_to_tensor('fac step:'),
                                       self.factor_step,
                                       tf.convert_to_tensor('sgd step:'),
                                       self.sgd_step,
                                       tf.convert_to_tensor('Accum:'),
                                       tf.convert_to_tensor(accumulate),
                                       tf.convert_to_tensor('Accum coeff:'),
                                       tf.convert_to_tensor(accumulateCoeff),
                                       tf.convert_to_tensor('stat step:'),
                                       self.stats_step, updateOps[0], updateOps[1]]))
        return [stats_step_op, ] 

def assign_add(ref, value):
  if hasattr(tf, "assign_add"):
    return tf.assign_add(ref, value)
  else:
    return ref.assign_add(value)