jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def _perform_update_steps(self, observ, action, old_mean, old_logstd, reward, length):
    """Perform multiple update steps of value function and policy.

    The advantage is computed once at the beginning and shared across
    iterations. We need to decide for the summary of one iteration, and thus
    choose the one after half of the iterations.

    Args:
      observ: Sequences of observations.
      action: Sequences of actions.
      old_mean: Sequences of action means of the behavioral policy.
      old_logstd: Sequences of action log stddevs of the behavioral policy.
      reward: Sequences of rewards.
      length: Batch of sequence lengths.

    Returns:
      Summary tensor.
    """
    return_ = utility.discounted_return(reward, length, self._config.discount)
    value = self._network(observ, length).value
    if self._config.gae_lambda:
      advantage = utility.lambda_return(reward, value, length, self._config.discount,
                                        self._config.gae_lambda)
    else:
      advantage = return_ - value
    mean, variance = tf.nn.moments(advantage, axes=[0, 1], keep_dims=True)
    advantage = (advantage - mean) / (tf.sqrt(variance) + 1e-8)
    advantage = tf.Print(advantage,
                         [tf.reduce_mean(return_), tf.reduce_mean(value)], 'return and value: ')
    advantage = tf.Print(advantage, [tf.reduce_mean(advantage)], 'normalized advantage: ')
    # pylint: disable=g-long-lambda
    value_loss, policy_loss, summary = tf.scan(lambda _1, _2: self._update_step(
        observ, action, old_mean, old_logstd, reward, advantage, length),
                                               tf.range(self._config.update_epochs), [0., 0., ''],
                                               parallel_iterations=1)
    print_losses = tf.group(tf.Print(0, [tf.reduce_mean(value_loss)], 'value loss: '),
                            tf.Print(0, [tf.reduce_mean(policy_loss)], 'policy loss: '))
    with tf.control_dependencies([value_loss, policy_loss, print_losses]):
      return summary[self._config.update_epochs // 2]
//  [python]
def _perform_update_steps(self, observ, action, old_mean, old_logstd, reward, length):
    """Perform multiple update steps of value function and policy.

    The advantage is computed once at the beginning and shared across
    iterations. We need to decide for the summary of one iteration, and thus
    choose the one after half of the iterations.

    Args:
      observ: Sequences of observations.
      action: Sequences of actions.
      old_mean: Sequences of action means of the behavioral policy.
      old_logstd: Sequences of action log stddevs of the behavioral policy.
      reward: Sequences of rewards.
      length: Batch of sequence lengths.

    Returns:
      Summary tensor.
    """
    return_ = utility.discounted_return(reward, length, self._config.discount)
    value = self._network(observ, length).value
    if self._config.gae_lambda:
      advantage = utility.lambda_return(reward, value, length, self._config.discount,
                                        self._config.gae_lambda)
    else:
      advantage = return_ - value
    mean, variance = tf.nn.moments(advantage, axes=[0, 1], keep_dims=True)
    advantage = (advantage - mean) / (tf.sqrt(variance) + 1e-8)
    advantage = tf.print(advantage,
                         [tf.reduce_mean(return_), tf.reduce_mean(value)], 'return and value: ')
    advantage = tf.print(advantage, [tf.reduce_mean(advantage)], 'normalized advantage: ')
    # pylint: disable=g-long-lambda
    value_loss, policy_loss, summary = tf.scan(lambda _1, _2: self._update_step(
        observ, action, old_mean, old_logstd, reward, advantage, length),
                                               tf.range(self._config.update_epochs), [0., 0., ''],
                                               parallel_iterations=1)
    print_losses = tf.group(tf.print(0, [tf.reduce_mean(value_loss)], 'value loss: '),
                            tf.print(0, [tf.reduce_mean(policy_loss)], 'policy loss: '))
    with tf.control_dependencies([value_loss, policy_loss, print_losses]):
      return summary[self._config.update_epochs // 2]
//  [python]
def debug_tensor(s, msg=None, summarize=10):
        """Print the shape and value of a tensor at test time. Return a new tensor."""
        if not msg:
                msg = s.name
        return tf.Print(s, [tf.shape(s), s], msg + " ", summarize=summarize)
//  [python]
def debug_tensor(s, msg=None, summarize=10):
        """Print the shape and value of a tensor at test time. Return a new tensor."""
        if not msg:
                msg = s.name
        return tf.print(s, [tf.shape(s), s], msg + " ", summarize=summarize)
//  [python]
def body(self, features, is_training):
    if is_training:
      init_features = features
      init_op = self.objective_tower(init_features, init=True)
      init_op = tf.Print(
          init_op, [init_op], message="Triggering data-dependent init.",
          first_n=20)
      tf.compat.v1.add_to_collection("glow_init_op", init_op)
    return self.objective_tower(features, init=False)
//  [python]
def body(self, features, is_training):
    if is_training:
      init_features = features
      init_op = self.objective_tower(init_features, init=True)
      init_op = tf.print(
          init_op, [init_op], message="Triggering data-dependent init.",
          first_n=20)
      tf.compat.v1.add_to_collection("glow_init_op", init_op)
    return self.objective_tower(features, init=False)
//  [python]
def print_act_stats(x, _str=""):
    if not do_print_act_stats:
        return x
    if hvd.rank() != 0:
        return x
    if len(x.get_shape()) == 1:
        x_mean, x_var = tf.nn.moments(x, [0], keep_dims=True)
    if len(x.get_shape()) == 2:
        x_mean, x_var = tf.nn.moments(x, [0], keep_dims=True)
    if len(x.get_shape()) == 4:
        x_mean, x_var = tf.nn.moments(x, [0, 1, 2], keep_dims=True)
    stats = [tf.reduce_min(x_mean), tf.reduce_mean(x_mean), tf.reduce_max(x_mean),
             tf.reduce_min(tf.sqrt(x_var)), tf.reduce_mean(tf.sqrt(x_var)), tf.reduce_max(tf.sqrt(x_var))]
    return tf.Print(x, stats, "["+_str+"] "+x.name)
//  [python]
def print_act_stats(x, _str=""):
    if not do_print_act_stats:
        return x
    if hvd.rank() != 0:
        return x
    if len(x.get_shape()) == 1:
        x_mean, x_var = tf.nn.moments(x, [0], keep_dims=True)
    if len(x.get_shape()) == 2:
        x_mean, x_var = tf.nn.moments(x, [0], keep_dims=True)
    if len(x.get_shape()) == 4:
        x_mean, x_var = tf.nn.moments(x, [0, 1, 2], keep_dims=True)
    stats = [tf.reduce_min(x_mean), tf.reduce_mean(x_mean), tf.reduce_max(x_mean),
             tf.reduce_min(tf.sqrt(x_var)), tf.reduce_mean(tf.sqrt(x_var)), tf.reduce_max(tf.sqrt(x_var))]
    return tf.print(x, stats, "["+_str+"] "+x.name)
//  [python]
def add_scalar_summary_op(tensor, name=None, 
    summary_key='summaries', print_summary_key='print_summaries', prefix=''):
  collections = []
  op = tf.summary.scalar(name, tensor, collections=collections)
  if summary_key != print_summary_key:
    tf.add_to_collection(summary_key, op)
  
  op = tf.Print(op, [tensor], '    {:-<25s}: '.format(name) + prefix)
  tf.add_to_collection(print_summary_key, op)
  return op
//  [python]
def add_scalar_summary_op(tensor, name=None, 
    summary_key='summaries', print_summary_key='print_summaries', prefix=''):
  collections = []
  op = tf.summary.scalar(name, tensor, collections=collections)
  if summary_key != print_summary_key:
    tf.add_to_collection(summary_key, op)
  
  op = tf.print(tensor, '    {:-<25s}: '.format(name) + prefix)
  tf.add_to_collection(print_summary_key, op)
  return op
//  [python]
def _max_pool(self, bottom, name, debug):
        pool = tf.nn.max_pool(bottom, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1],
                              padding='SAME', name=name)

        if debug:
            pool = tf.Print(pool, [tf.shape(pool)],
                            message='Shape of %s' % name,
                            summarize=4, first_n=1)
        return pool
//  [python]
def _max_pool(self, bottom, name, debug):
        pool = tf.nn.max_pool(bottom, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1],
                              padding='SAME', name=name)

        if debug:
            pool = tf.print(pool, [tf.shape(pool)],
                            message='Shape of %s' % name,
                            summarize=4, first_n=1)
        return pool
//  [python]
def reduce_hitrate(cluster, hits, count, task_index):
  """Reduce hitrate of all workers.

  Args:
    cluster: tf cluster.
    hits: total_hits of each worker.
    count: total count of ground truth items of each worker.
    task_index: worker index.

  Returns:
    var_total_hitrate: variable of total hitrate.
    var_worker_count: variable used to mark the number of worker that
    have completed the calculation of hitrate.
  """
  with tf.device(
      tf.train.replica_device_setter(
          worker_device='/job:worker/task:%d' % task_index, cluster=cluster)):
    with tf.variable_scope('hitrate_var', reuse=tf.AUTO_REUSE):
      var_worker_count = tf.get_variable(
          'worker_count',
          shape=(),
          dtype=tf.int32,
          initializer=tf.zeros_initializer())
      var_hits = tf.get_variable(
          'hits',
          shape=(),
          dtype=tf.float32,
          initializer=tf.zeros_initializer())
      var_gt_count = tf.get_variable(
          'gt_count',
          shape=(),
          dtype=tf.float32,
          initializer=tf.zeros_initializer())
      var_total_hitrate = tf.get_variable(
          'total_hitate',
          shape=(),
          dtype=tf.float32,
          initializer=tf.zeros_initializer())

      var_hits = tf.assign_add(var_hits, hits, use_locking=True)
      var_gt_count = tf.assign_add(var_gt_count, count, use_locking=True)
      var_gt_count = tf.Print(
          var_gt_count, [var_gt_count, var_hits],
          message='var_gt_count/var_hits')
      var_total_hitrate = tf.assign(
          var_total_hitrate, var_hits / var_gt_count, use_locking=True)
      with tf.control_dependencies([var_total_hitrate]):
        var_worker_count = tf.assign_add(var_worker_count, 1, use_locking=True)
  return var_total_hitrate, var_worker_count
//  [python]
def reduce_hitrate(cluster, hits, count, task_index):
  """Reduce hitrate of all workers.

  Args:
    cluster: tf cluster.
    hits: total_hits of each worker.
    count: total count of ground truth items of each worker.
    task_index: worker index.

  Returns:
    var_total_hitrate: variable of total hitrate.
    var_worker_count: variable used to mark the number of worker that
    have completed the calculation of hitrate.
  """
  with tf.device(
      tf.train.replica_device_setter(
          worker_device='/job:worker/task:%d' % task_index, cluster=cluster)):
    with tf.variable_scope('hitrate_var', reuse=tf.AUTO_REUSE):
      var_worker_count = tf.get_variable(
          'worker_count',
          shape=(),
          dtype=tf.int32,
          initializer=tf.zeros_initializer())
      var_hits = tf.get_variable(
          'hits',
          shape=(),
          dtype=tf.float32,
          initializer=tf.zeros_initializer())
      var_gt_count = tf.get_variable(
          'gt_count',
          shape=(),
          dtype=tf.float32,
          initializer=tf.zeros_initializer())
      var_total_hitrate = tf.get_variable(
          'total_hitate',
          shape=(),
          dtype=tf.float32,
          initializer=tf.zeros_initializer())

      var_hits = tf.assign_add(var_hits, hits, use_locking=True)
      var_gt_count = tf.assign_add(var_gt_count, count, use_locking=True)
      var_gt_count = tf.print(
          var_gt_count, [var_gt_count, var_hits],
          message='var_gt_count/var_hits')
      var_total_hitrate = tf.assign(
          var_total_hitrate, var_hits / var_gt_count, use_locking=True)
      with tf.control_dependencies([var_total_hitrate]):
        var_worker_count = tf.assign_add(var_worker_count, 1, use_locking=True)
  return var_total_hitrate, var_worker_count
//  [python]
def detectMinVal(input_mat, var, threshold=1e-6, name='', debug=False):
    eigen_min = tf.reduce_min(input_mat)
    eigen_max = tf.reduce_max(input_mat)
    eigen_ratio = eigen_max / eigen_min
    input_mat_clipped = clipoutNeg(input_mat, threshold)

    if debug:
        input_mat_clipped = tf.cond(tf.logical_or(tf.greater(eigen_ratio, 0.), tf.less(eigen_ratio, -500)), lambda: input_mat_clipped, lambda: tf.Print(
            input_mat_clipped, [tf.convert_to_tensor('screwed ratio ' + name + ' eigen values!!!'), tf.convert_to_tensor(var.name), eigen_min, eigen_max, eigen_ratio]))

    return input_mat_clipped
//  [python]
def detectMinVal(input_mat, var, threshold=1e-6, name='', debug=False):
    eigen_min = tf.reduce_min(input_mat)
    eigen_max = tf.reduce_max(input_mat)
    eigen_ratio = eigen_max / eigen_min
    input_mat_clipped = clipoutNeg(input_mat, threshold)

    if debug:
        input_mat_clipped = tf.cond(tf.logical_or(tf.greater(eigen_ratio, 0.), tf.less(eigen_ratio, -500)), lambda: input_mat_clipped, lambda: tf.print(
            input_mat_clipped, ['screwed ratio ' + name + ' eigen values!!!', var.name, eigen_min, eigen_max, eigen_ratio]))

    return input_mat_clipped
//  [python]
def _fc_layer(self, bottom, name, num_classes=None,
                  relu=True, debug=False):
        with tf.variable_scope(name) as scope:
            shape = bottom.get_shape().as_list()

            if name == 'fc6':
                filt = self.get_fc_weight_reshape(name, [7, 7, 512, 4096])
            elif name == 'score_fr':
                name = 'fc8'  # Name of score_fr layer in VGG Model
                filt = self.get_fc_weight_reshape(name, [1, 1, 4096, 1000],
                                                  num_classes=num_classes)
            else:
                filt = self.get_fc_weight_reshape(name, [1, 1, 4096, 4096])

            self._add_wd_and_summary(filt, self.wd, "fc_wlosses")

            conv = tf.nn.conv2d(bottom, filt, [1, 1, 1, 1], padding='SAME')
            conv_biases = self.get_bias(name, num_classes=num_classes)
            bias = tf.nn.bias_add(conv, conv_biases)

            if relu:
                bias = tf.nn.relu(bias)
            _activation_summary(bias)

            if debug:
                bias = tf.Print(bias, [tf.shape(bias)],
                                message='Shape of %s' % name,
                                summarize=4, first_n=1)
            return bias
//  [python]
def _fc_layer(self, bottom, name, num_classes=None,
                  relu=True, debug=False):
        with tf.variable_scope(name) as scope:
            shape = bottom.get_shape().as_list()

            if name == 'fc6':
                filt = self.get_fc_weight_reshape(name, [7, 7, 512, 4096])
            elif name == 'score_fr':
                name = 'fc8'  # Name of score_fr layer in VGG Model
                filt = self.get_fc_weight_reshape(name, [1, 1, 4096, 1000],
                                                  num_classes=num_classes)
            else:
                filt = self.get_fc_weight_reshape(name, [1, 1, 4096, 4096])

            self._add_wd_and_summary(filt, self.wd, "fc_wlosses")

            conv = tf.nn.conv2d(bottom, filt, [1, 1, 1, 1], padding='SAME')
            conv_biases = self.get_bias(name, num_classes=num_classes)
            bias = tf.nn.bias_add(conv, conv_biases)

            if relu:
                bias = tf.nn.relu(bias)
            _activation_summary(bias)

            if debug:
                bias = tf.print(bias, [tf.shape(bias)],
                                message='Shape of %s' % name,
                                summarize=4, first_n=1)
            return bias
//  [python]
def build(self, predictions, targets, inputs=None):
        """ Prints the number of each kind of prediction """
        self.built = True
        pshape = predictions.get_shape()
        self.inner_metric.build(predictions, targets, inputs)

        with tf.name_scope(self.name):
            if len(pshape) == 1 or (len(pshape) == 2 and int(pshape[1]) == 1):
                self.name = self.name or "binary_prediction_counts"
                y, idx, count = tf.unique_with_counts(tf.argmax(predictions))
                self.tensor = tf.Print(self.inner_metric, [y, count], name=self.inner_metric.name)
            else:
                self.name = self.name or "categorical_prediction_counts"
                y, idx, count = tf.unique_with_counts(tf.argmax(predictions, dimension=1))
                self.tensor = tf.Print(self.inner_metric.tensor, [y, count], name=self.inner_metric.name)
//  [python]
def build(self, predictions, targets, inputs=None):
        """ Prints the number of each kind of prediction """
        self.built = True
        pshape = predictions.get_shape()
        self.inner_metric.build(predictions, targets, inputs)

        with tf.name_scope(self.name):
            if len(pshape) == 1 or (len(pshape) == 2 and int(pshape[1]) == 1):
                self.name = self.name or "binary_prediction_counts"
                y, idx, count = tf.unique_with_counts(tf.argmax(predictions))
                self.tensor = tf.print(self.inner_metric, [y, count], name=self.inner_metric.name)
            else:
                self.name = self.name or "categorical_prediction_counts"
                y, idx, count = tf.unique_with_counts(tf.argmax(predictions, dimension=1))
                self.tensor = tf.print(self.inner_metric.tensor, [y, count], name=self.inner_metric.name)
//  [python]
def _upscore_layer(self, bottom, shape,
                       num_classes, name, debug,
                       ksize=4, stride=2):
        strides = [1, stride, stride, 1]
        with tf.variable_scope(name):
            in_features = bottom.get_shape()[3].value

            if shape is None:
                # Compute shape out of Bottom
                in_shape = tf.shape(bottom)

                h = ((in_shape[1] - 1) * stride) + 1
                w = ((in_shape[2] - 1) * stride) + 1
                new_shape = [in_shape[0], h, w, num_classes]
            else:
                new_shape = [shape[0], shape[1], shape[2], num_classes]
            output_shape = tf.stack(new_shape)

            logging.debug("Layer: %s, Fan-in: %d" % (name, in_features))
            f_shape = [ksize, ksize, num_classes, in_features]

            # create
            num_input = ksize * ksize * in_features / stride
            stddev = (2 / num_input)**0.5

            weights = self.get_deconv_filter(f_shape)
            self._add_wd_and_summary(weights, self.wd, "fc_wlosses")
            deconv = tf.nn.conv2d_transpose(bottom, weights, output_shape,
                                            strides=strides, padding='SAME')

            if debug:
                deconv = tf.Print(deconv, [tf.shape(deconv)],
                                  message='Shape of %s' % name,
                                  summarize=4, first_n=1)

        _activation_summary(deconv)
        return deconv
//  [python]
def _upscore_layer(self, bottom, shape,
                       num_classes, name, debug,
                       ksize=4, stride=2):
        strides = [1, stride, stride, 1]
        with tf.variable_scope(name):
            in_features = bottom.get_shape()[3].value

            if shape is None:
                # Compute shape out of Bottom
                in_shape = tf.shape(bottom)

                h = ((in_shape[1] - 1) * stride) + 1
                w = ((in_shape[2] - 1) * stride) + 1
                new_shape = [in_shape[0], h, w, num_classes]
            else:
                new_shape = [shape[0], shape[1], shape[2], num_classes]
            output_shape = tf.stack(new_shape)

            logging.debug("Layer: %s, Fan-in: %d" % (name, in_features))
            f_shape = [ksize, ksize, num_classes, in_features]

            # create
            num_input = ksize * ksize * in_features / stride
            stddev = (2 / num_input)**0.5

            weights = self.get_deconv_filter(f_shape)
            self._add_wd_and_summary(weights, self.wd, "fc_wlosses")
            deconv = tf.nn.conv2d_transpose(bottom, weights, output_shape,
                                            strides=strides, padding='SAME')

            if debug:
                deconv = tf.print(deconv, [tf.shape(deconv)],
                                  message='Shape of %s' % name,
                                  summarize=4, first_n=1)

        _activation_summary(deconv)
        return deconv
//  [python]
def build_evaluation_graph(loss_module_fn=gin.REQUIRED,
                           learner_fn=gin.REQUIRED):
  """Build the evaluation graph for inner-training."""
  global_step = tf.train.get_or_create_global_step()

  loss_module = loss_module_fn()
  learner, theta_mod = learner_fn(loss_module)

  initial_state = learner.initial_state()
  reset_state_op = learner.assign_state(initial_state)
  state = learner.current_state()
  state = tf_utils.force_copy(state)

  with tf.control_dependencies(nest.flatten(state)):
    last_loss, new_state = learner.loss_and_next_state(state)

  with tf.control_dependencies([last_loss] + nest.flatten(new_state)):
    train_op = learner.assign_state(new_state)

  update_global_step = global_step.assign_add(1)
  train_op = tf.group(train_op, update_global_step, name="train_op")

  load_vars = list(theta_mod.get_variables(tf.GraphKeys.GLOBAL_VARIABLES))
  meta_loss = learner.meta_loss(state)
  meta_loss = tf.Print(meta_loss, [meta_loss], "meta_loss")

  inner_loss = learner.inner_loss(state)
  inner_loss = tf.Print(inner_loss, [inner_loss], "inner_loss")
  # TODO(lmetz) this should only be scalars.
  train_op = tf.group(train_op, name="train_op")

  return {
      "train_op": train_op,
      "global_step": global_step,
      "init_op": reset_state_op,
      "checkpoint_vars": load_vars,
      "meta_loss": meta_loss,
      "train_loss": inner_loss,
      "current_state": state,
      "next_state": new_state,
  }
//  [python]
def build_evaluation_graph(loss_module_fn=gin.REQUIRED,
                           learner_fn=gin.REQUIRED):
  """Build the evaluation graph for inner-training."""
  global_step = tf.train.get_or_create_global_step()

  loss_module = loss_module_fn()
  learner, theta_mod = learner_fn(loss_module)

  initial_state = learner.initial_state()
  reset_state_op = learner.assign_state(initial_state)
  state = learner.current_state()
  state = tf_utils.force_copy(state)

  with tf.control_dependencies(nest.flatten(state)):
    last_loss, new_state = learner.loss_and_next_state(state)

  with tf.control_dependencies([last_loss] + nest.flatten(new_state)):
    train_op = learner.assign_state(new_state)

  update_global_step = global_step.assign_add(1)
  train_op = tf.group(train_op, update_global_step, name="train_op")

  load_vars = list(theta_mod.get_variables(tf.GraphKeys.GLOBAL_VARIABLES))
  meta_loss = learner.meta_loss(state)
  meta_loss = tf.print(meta_loss, [meta_loss], "meta_loss")

  inner_loss = learner.inner_loss(state)
  inner_loss = tf.print(inner_loss, [inner_loss], "inner_loss")
  # TODO(lmetz) this should only be scalars.
  train_op = tf.group(train_op, name="train_op")

  return {
      "train_op": train_op,
      "global_step": global_step,
      "init_op": reset_state_op,
      "checkpoint_vars": load_vars,
      "meta_loss": meta_loss,
      "train_loss": inner_loss,
      "current_state": state,
      "next_state": new_state,
  }
//  [python]
def benchmark_handwritten(self):
    with tf.Graph().as_default():
      ds, opt, hp, w, b = get_data_and_params()
      iterator = ds.make_one_shot_iterator()

      def loop_body(i, unused_previous_loss_t):
        """Manual implementation of training loop."""
        # Call get_next() inside body or else training happens repeatedly on
        # the first minibatch only.
        x, y = iterator.get_next()
        loss_t = loss_fn(x, y, w, b)
        train_op = opt.minimize(loss_t, var_list=(w, b))
        i = tf.cond(tf.equal(i % 100, 0),
                    lambda: tf.Print(i, [i, loss_t], message='Step, loss: '),
                    lambda: i)

        with tf.control_dependencies([train_op]):
          return i + 1, loss_t

      _, final_loss_t = tf.while_loop(
          lambda i, _: i < hp.train_steps,
          loop_body,
          [tf.constant(0), tf.constant(0.0)])

      with tf.Session() as sess:
        sess.run(tf.global_variables_initializer())

        def target():
          loss_val = sess.run(final_loss_t)
          assert 0.1 < loss_val < 1, loss_val

        self.time_execution(
            'Handwritten',
            target,
            iter_volume=hp.train_steps,
            iter_unit='training steps')
//  [python]
def benchmark_handwritten(self):
    with tf.Graph().as_default():
      ds, opt, hp, w, b = get_data_and_params()
      iterator = ds.make_one_shot_iterator()

      def loop_body(i, unused_previous_loss_t):
        """Manual implementation of training loop."""
        # Call get_next() inside body or else training happens repeatedly on
        # the first minibatch only.
        x, y = iterator.get_next()
        loss_t = loss_fn(x, y, w, b)
        train_op = opt.minimize(loss_t, var_list=(w, b))
        i = tf.cond(tf.equal(i % 100, 0),
                    lambda: tf.print(i, [i, loss_t], message='Step, loss: '),
                    lambda: i)

        with tf.control_dependencies([train_op]):
          return i + 1, loss_t

      _, final_loss_t = tf.while_loop(
          lambda i, _: i < hp.train_steps,
          loop_body,
          [tf.constant(0), tf.constant(0.0)])

      with tf.Session() as sess:
        sess.run(tf.global_variables_initializer())

        def target():
          loss_val = sess.run(final_loss_t)
          assert 0.1 < loss_val < 1, loss_val

        self.time_execution(
            'Handwritten',
            target,
            iter_volume=hp.train_steps,
            iter_unit='training steps')
//  [python]
def print_data_example(logits, lolly_activations, features):
  return tf.Print(
    logits,
    [logits, lolly_activations, tf.reshape(features['keys'], (1, -1)), tf.reshape(tf.multiply(features['values'], -1.0), (1, -1))],
    message="DATA EXAMPLE = ",
    summarize=10000
  )
//  [python]
def print_data_example(logits, lolly_activations, features):
  return tf.print(
    logits,
    [logits, lolly_activations, tf.reshape(features['keys'], (1, -1)), tf.reshape(tf.multiply(features['values'], -1.0), (1, -1))],
    message="DATA EXAMPLE = ",
    summarize=10000
  )
//  [python]
def print_debug_statistics(image, labels, dbg_messge_prefix, on_tpu):
  """Adds a Print directive to an image tensor which prints debug statistics."""
  if on_tpu:
    # Print operations are not supported on TPUs.
    return image, labels
  image_means = tf.reduce_mean(input_tensor=image, axis=0, keepdims=True)
  image_vars = tf.reduce_mean(
      input_tensor=tf.math.squared_difference(image, image_means),
      axis=0,
      keepdims=True)
  image = tf.Print(
      image, [
          tf.reduce_mean(input_tensor=image_means),
          tf.reduce_mean(input_tensor=image_vars)
      ],
      dbg_messge_prefix + ' mean and average var',
      first_n=1)
  labels = tf.Print(
      labels, [labels, labels.shape],
      dbg_messge_prefix + ' sparse labels',
      first_n=2)
  return image, labels
//  [python]
def print_debug_statistics(image, labels, dbg_messge_prefix, on_tpu):
  """Adds a Print directive to an image tensor which prints debug statistics."""
  if on_tpu:
    # Print operations are not supported on TPUs.
    return image, labels
  image_means = tf.reduce_mean(input_tensor=image, axis=0, keepdims=True)
  image_vars = tf.reduce_mean(
      input_tensor=tf.math.squared_difference(image, image_means),
      axis=0,
      keepdims=True)
  image = tf.print(
      image, [
          tf.reduce_mean(input_tensor=image_means),
          tf.reduce_mean(input_tensor=image_vars)
      ],
      dbg_messge_prefix + ' mean and average var',
      first_n=1)
  labels = tf.print(
      labels, [labels, labels.shape],
      dbg_messge_prefix + ' sparse labels',
      first_n=2)
  return image, labels
//  [python]
def build_graph_train(self, x_l, y_l, x, is_supervised=True):

        o = dict()  # output
        loss = 0

        logit = self.encoder(x)

        with tf.variable_scope(tf.get_variable_scope(), reuse=True):
            logit_l = self.encoder(x_l, is_train=True, do_update_bn=False)  # for pyx and vat loss computation

        """ Classification Loss """
        o['Ly'], o['accur'] = self.lf.get_loss_pyx(logit_l, y_l)
        loss += o['Ly']

        """ PI Model Loss """
        if self.use_pi:
            with tf.variable_scope(tf.get_variable_scope(), reuse=True):
                _,_,o['Lp'] = self.lf.get_loss_pi(x, logit, is_train=True)
                loss += self.lambda_pi_usl * o['Lp']
        else:
            o['Lp'] = tf.constant(0)

        """ set losses """
        o['loss'] = loss
        self.o_train = o

        """ set optimizer """
        optimizer = tf.train.AdamOptimizer(learning_rate=self.lr, beta1=0.5)
        #self.op = optimizer.minimize(loss)
        grads = optimizer.compute_gradients(loss)
        for i,(g,v) in enumerate(grads):
            if g is not None:
                #g = tf.Print(g, [g], "g %s = "%(v))
                grads[i] = (tf.clip_by_norm(g,5),v) # clip gradients
            else:
                print('g is None:', v)
                v = tf.Print(v, [v], "v = ", summarize=10000)
        self.op = optimizer.apply_gradients(grads)
//  [python]
def build_graph_train(self, x_l, y_l, x, is_supervised=True):

        o = dict()  # output
        loss = 0

        logit = self.encoder(x)

        with tf.variable_scope(tf.get_variable_scope(), reuse=True):
            logit_l = self.encoder(x_l, is_train=True, do_update_bn=False)  # for pyx and vat loss computation

        """ Classification Loss """
        o['Ly'], o['accur'] = self.lf.get_loss_pyx(logit_l, y_l)
        loss += o['Ly']

        """ PI Model Loss """
        if self.use_pi:
            with tf.variable_scope(tf.get_variable_scope(), reuse=True):
                _,_,o['Lp'] = self.lf.get_loss_pi(x, logit, is_train=True)
                loss += self.lambda_pi_usl * o['Lp']
        else:
            o['Lp'] = tf.constant(0)

        """ set losses """
        o['loss'] = loss
        self.o_train = o

        """ set optimizer """
        optimizer = tf.train.AdamOptimizer(learning_rate=self.lr, beta1=0.5)
        #self.op = optimizer.minimize(loss)
        grads = optimizer.compute_gradients(loss)
        for i,(g,v) in enumerate(grads):
            if g is not None:
                #g = tf.print(g, [g], "g %s = "%(v))
                grads[i] = (tf.clip_by_norm(g,5),v) # clip gradients
            else:
                print('g is None:', v)
                v = tf.print(v, [v], "v = ", summarize=10000)
        self.op = optimizer.apply_gradients(grads)
//  [python]
def debug(x, idx):
  return tf.Print(x,
                  [idx, tf.reduce_mean(x), tf.reduce_max(x), tf.reduce_min(x)])
//  [python]
def debug(x, idx):
  return tf.print(idx, tf.reduce_mean(x), tf.reduce_max(x), tf.reduce_min(x))
//  [python]
def compute_safe_IOU(target_rois, detected_rois, detected_rois_overflow, tile_size):
    """Computes the Intersection Over Union (IOU) of a batch of detected boxes
    against a batch of target boxes. Logs a message if a problem occurs."""

    iou_accuracy = IOUCalculator.batch_intersection_over_union(detected_rois * tile_size, target_rois * tile_size, tile_size=tile_size)
    iou_accuracy_overflow = tf.greater(tf.reduce_sum(detected_rois_overflow), 0)
    # check that we are not overflowing the tensor size. Issue a warning if we are. This should only happen at
    # the beginning of the training with a completely uninitialized network.
    iou_accuracy = tf.cond(iou_accuracy_overflow,
                           lambda: tf.Print(iou_accuracy, [detected_rois_overflow],
                                            summarize=250, message="ROI tensor overflow in IOU computation. "
                                                                   "The computed IOU is not correct and will "
                                                                   "be reported as 0. This can be normal in initial "
                                                                   "training iteration when all weights are random. "
                                                                   "Increase MAX_DETECTED_ROIS_PER_TILE to avoid."),
                           lambda: tf.identity(iou_accuracy))
    iou_accuracy = IOUCalculator.batch_mean(iou_accuracy)
    # set iou_accuracy to 0 if there has been any overflow in its computation
    iou_accuracy = tf.where(iou_accuracy_overflow, tf.zeros_like(iou_accuracy), iou_accuracy)
    return iou_accuracy
//  [python]
def compute_safe_IOU(target_rois, detected_rois, detected_rois_overflow, tile_size):
    """Computes the Intersection Over Union (IOU) of a batch of detected boxes
    against a batch of target boxes. Logs a message if a problem occurs."""

    iou_accuracy = IOUCalculator.batch_intersection_over_union(detected_rois * tile_size, target_rois * tile_size, tile_size=tile_size)
    iou_accuracy_overflow = tf.greater(tf.reduce_sum(detected_rois_overflow), 0)
    # check that we are not overflowing the tensor size. Issue a warning if we are. This should only happen at
    # the beginning of the training with a completely uninitialized network.
    iou_accuracy = tf.cond(iou_accuracy_overflow,
                           lambda: tf.print(iou_accuracy, [detected_rois_overflow],
                                            summarize=250, message="ROI tensor overflow in IOU computation. "
                                                                   "The computed IOU is not correct and will "
                                                                   "be reported as 0. This can be normal in initial "
                                                                   "training iteration when all weights are random. "
                                                                   "Increase MAX_DETECTED_ROIS_PER_TILE to avoid."),
                           lambda: tf.identity(iou_accuracy))
    iou_accuracy = IOUCalculator.batch_mean(iou_accuracy)
    # set iou_accuracy to 0 if there has been any overflow in its computation
    iou_accuracy = tf.where(iou_accuracy_overflow, tf.zeros_like(iou_accuracy), iou_accuracy)
    return iou_accuracy
//  [python]
def reorder_beam(beam_size, batch_size, beam_val, output, is_first,
                 tensors_to_reorder):
  """Reorder to minimize beam costs."""
  # beam_val is [batch_size x beam_size]; let b = batch_size * beam_size
  # decided is len x b x a x b
  # output is b x out_size; step is b x len x a x b;
  outputs = tf.split(axis=0, num_or_size_splits=beam_size, value=tf.nn.log_softmax(output))
  all_beam_vals, all_beam_idx = [], []
  beam_range = 1 if is_first else beam_size
  for i in xrange(beam_range):
    top_out, top_out_idx = tf.nn.top_k(outputs[i], k=beam_size)
    cur_beam_val = beam_val[:, i]
    top_out = tf.Print(top_out, [top_out, top_out_idx, beam_val, i,
                                 cur_beam_val], "GREPO", summarize=8)
    all_beam_vals.append(top_out + tf.expand_dims(cur_beam_val, 1))
    all_beam_idx.append(top_out_idx)
  all_beam_idx = tf.reshape(tf.transpose(tf.concat(axis=1, values=all_beam_idx), [1, 0]),
                            [-1])
  top_beam, top_beam_idx = tf.nn.top_k(tf.concat(axis=1, values=all_beam_vals), k=beam_size)
  top_beam_idx = tf.Print(top_beam_idx, [top_beam, top_beam_idx],
                          "GREP", summarize=8)
  reordered = [[] for _ in xrange(len(tensors_to_reorder) + 1)]
  top_out_idx = []
  for i in xrange(beam_size):
    which_idx = top_beam_idx[:, i] * batch_size + tf.range(batch_size)
    top_out_idx.append(tf.gather(all_beam_idx, which_idx))
    which_beam = top_beam_idx[:, i] / beam_size  # [batch]
    which_beam = which_beam * batch_size + tf.range(batch_size)
    reordered[0].append(tf.gather(output, which_beam))
    for i, t in enumerate(tensors_to_reorder):
      reordered[i + 1].append(tf.gather(t, which_beam))
  new_tensors = [tf.concat(axis=0, values=t) for t in reordered]
  top_out_idx = tf.concat(axis=0, values=top_out_idx)
  return (top_beam, new_tensors[0], top_out_idx, new_tensors[1:])
//  [python]
def reorder_beam(beam_size, batch_size, beam_val, output, is_first,
                 tensors_to_reorder):
  """Reorder to minimize beam costs."""
  # beam_val is [batch_size x beam_size]; let b = batch_size * beam_size
  # decided is len x b x a x b
  # output is b x out_size; step is b x len x a x b;
  outputs = tf.split(axis=0, num_or_size_splits=beam_size, value=tf.nn.log_softmax(output))
  all_beam_vals, all_beam_idx = [], []
  beam_range = 1 if is_first else beam_size
  for i in xrange(beam_range):
    top_out, top_out_idx = tf.nn.top_k(outputs[i], k=beam_size)
    cur_beam_val = beam_val[:, i]
    top_out = tf.print(top_out, [top_out, top_out_idx, beam_val, i,
                                 cur_beam_val], "GREPO", summarize=8)
    all_beam_vals.append(top_out + tf.expand_dims(cur_beam_val, 1))
    all_beam_idx.append(top_out_idx)
  all_beam_idx = tf.reshape(tf.transpose(tf.concat(axis=1, values=all_beam_idx), [1, 0]),
                            [-1])
  top_beam, top_beam_idx = tf.nn.top_k(tf.concat(axis=1, values=all_beam_vals), k=beam_size)
  top_beam_idx = tf.print(top_beam_idx, [top_beam, top_beam_idx],
                          "GREP", summarize=8)
  reordered = [[] for _ in xrange(len(tensors_to_reorder) + 1)]
  top_out_idx = []
  for i in xrange(beam_size):
    which_idx = top_beam_idx[:, i] * batch_size + tf.range(batch_size)
    top_out_idx.append(tf.gather(all_beam_idx, which_idx))
    which_beam = top_beam_idx[:, i] / beam_size  # [batch]
    which_beam = which_beam * batch_size + tf.range(batch_size)
    reordered[0].append(tf.gather(output, which_beam))
    for i, t in enumerate(tensors_to_reorder):
      reordered[i + 1].append(tf.gather(t, which_beam))
  new_tensors = [tf.concat(axis=0, values=t) for t in reordered]
  top_out_idx = tf.concat(axis=0, values=top_out_idx)
  return (top_beam, new_tensors[0], top_out_idx, new_tensors[1:])
//  [python]
def prune(signals, masks=None):
  """
  Prunes a single feature by zeroing the mask entry with the smallest pruning signal.

  Arguments:
    signals: list[tf.Tensor]
      A list of pruning signals

    masks: list[tf.Tensor]
      A list of corresponding masks, defaults to `tf.get_collection(MASK_COLLECTION)`

  Returns:
    A `tf.Operation` which updates masks
  """

  if masks is None:
    masks = tf.get_collection(MASK_COLLECTION)

  with tf.variable_scope('pruning_opt', reuse=True):
    # make sure we don't select already pruned units
    signals = [tf.where(m > .5, s, tf.zeros_like(s) + np.inf) for m, s in zip(masks, signals)]

    # find units with smallest pruning signal in each layer
    min_idx = [tf.argmin(s) for s in signals]
    min_signals = [s[i] for s, i in zip(signals, min_idx)]

    # find layer with smallest pruning signal
    l = tf.argmin(min_signals)

    # construct pruning operations, one for each mask
    updates = []
    for k, i in enumerate(min_idx):
      # set mask of layer l to 0 where pruning signal is smallest
      updates.append(
        tf.cond(
          tf.equal(l, k),
          lambda: tf.scatter_update(
            masks[k], tf.Print(i, [i], message="Pruning layer [{0}] at index ".format(k)), 0.),
          lambda: masks[k]))

    updates = tf.group(updates, name='prune')

  return updates
//  [python]
def prune(signals, masks=None):
  """
  Prunes a single feature by zeroing the mask entry with the smallest pruning signal.

  Arguments:
    signals: list[tf.Tensor]
      A list of pruning signals

    masks: list[tf.Tensor]
      A list of corresponding masks, defaults to `tf.get_collection(MASK_COLLECTION)`

  Returns:
    A `tf.Operation` which updates masks
  """

  if masks is None:
    masks = tf.get_collection(MASK_COLLECTION)

  with tf.variable_scope('pruning_opt', reuse=True):
    # make sure we don't select already pruned units
    signals = [tf.where(m > .5, s, tf.zeros_like(s) + np.inf) for m, s in zip(masks, signals)]

    # find units with smallest pruning signal in each layer
    min_idx = [tf.argmin(s) for s in signals]
    min_signals = [s[i] for s, i in zip(signals, min_idx)]

    # find layer with smallest pruning signal
    l = tf.argmin(min_signals)

    # construct pruning operations, one for each mask
    updates = []
    for k, i in enumerate(min_idx):
      # set mask of layer l to 0 where pruning signal is smallest
      updates.append(
        tf.cond(
          tf.equal(l, k),
          lambda: tf.scatter_update(
            masks[k], tf.print(i, message="Pruning layer [{0}] at index ".format(k)), 0.),
          lambda: masks[k]))

    updates = tf.group(updates, name='prune')

  return updates
//  [python]
def print_abs_mean(node, tensor, info = ""):
  return tf.Print(node, ["[INFO] {} abs mean".format(info),
                         tf.shape(tensor),
                         tf.reduce_mean(tf.abs(tensor)),
                         tf.reduce_sum(tf.abs(tensor))])
//  [python]
def print_abs_mean(node, tensor, info = ""):
  return tf.print("[INFO] {} abs mean".format(info),
                         tf.shape(tensor),
                         tf.reduce_mean(tf.abs(tensor)),
                         tf.reduce_sum(tf.abs(tensor)),
                         output_stream=sys.stdout)
//  [python]
def __init__(self, filters, blocks, components, heads=4, init_scale=0.1, enable_print=True):
        def f(x, *, vcfg: VarConfig, context=None, dropout_p=0., verbose=True):
            if vcfg.init and verbose and enable_print:
                # debug stuff
                xmean, xvar = tf.nn.moments(x, axes=list(range(len(x.shape))))
                x = tf.Print(
                    x, [tf.shape(x), xmean, tf.sqrt(xvar), tf.reduce_min(x), tf.reduce_max(x)],
                    message='{} (shape/mean/std/min/max) '.format(self.template.variable_scope.name), summarize=10
                )
            B, H, W, C = x.shape.as_list()
            pos_emb = get_var('pos_emb', shape=[H, W, filters], initializer=tf.random_normal_initializer(stddev=0.01),
                              vcfg=vcfg)
            x = conv2d(x, name='proj_in', num_units=filters, vcfg=vcfg)
            for i_block in range(blocks):
                with tf.variable_scope('block{}'.format(i_block)):
                    x = gated_conv(x, name='conv', a=context, use_nin=True, dropout_p=dropout_p, vcfg=vcfg)
                    x = layernorm(x, name='ln1', vcfg=vcfg)
                    x = gated_attn(x, name='attn', pos_emb=pos_emb, heads=heads, dropout_p=dropout_p, vcfg=vcfg)
                    x = layernorm(x, name='ln2', vcfg=vcfg)
            x = conv2d(x, name='proj_out', num_units=C * (2 + 3 * components), init_scale=init_scale, vcfg=vcfg)
            assert x.shape == [B, H, W, C * (2 + 3 * components)]
            x = tf.reshape(x, [B, H, W, C, 2 + 3 * components])

            s, t = tf.tanh(x[:, :, :, :, 0]), x[:, :, :, :, 1]
            ml_logits, ml_means, ml_logscales = tf.split(x[:, :, :, :, 2:], 3, axis=4)
            assert s.shape == t.shape == [B, H, W, C]
            assert ml_logits.shape == ml_means.shape == ml_logscales.shape == [B, H, W, C, components]

            return Compose([
                MixLogisticCDF(logits=ml_logits, means=ml_means, logscales=ml_logscales),
                Inverse(Sigmoid()),
                ElemwiseAffine(scales=tf.exp(s), logscales=s, biases=t),
            ])

        self.template = tf.make_template(self.__class__.__name__, f)
//  [python]
def __init__(self, filters, blocks, components, heads=4, init_scale=0.1, enable_print=True):
        def f(x, *, vcfg: VarConfig, context=None, dropout_p=0., verbose=True):
            if vcfg.init and verbose and enable_print:
                # debug stuff
                xmean, xvar = tf.nn.moments(x, axes=list(range(len(x.shape))))
                x = tf.print(
                    x, [tf.shape(x), xmean, tf.sqrt(xvar), tf.reduce_min(x), tf.reduce_max(x)],
                    message='{} (shape/mean/std/min/max) '.format(self.template.variable_scope.name), summarize=10
                )
            B, H, W, C = x.shape.as_list()
            pos_emb = get_var('pos_emb', shape=[H, W, filters], initializer=tf.random_normal_initializer(stddev=0.01),
                              vcfg=vcfg)
            x = conv2d(x, name='proj_in', num_units=filters, vcfg=vcfg)
            for i_block in range(blocks):
                with tf.variable_scope('block{}'.format(i_block)):
                    x = gated_conv(x, name='conv', a=context, use_nin=True, dropout_p=dropout_p, vcfg=vcfg)
                    x = layernorm(x, name='ln1', vcfg=vcfg)
                    x = gated_attn(x, name='attn', pos_emb=pos_emb, heads=heads, dropout_p=dropout_p, vcfg=vcfg)
                    x = layernorm(x, name='ln2', vcfg=vcfg)
            x = conv2d(x, name='proj_out', num_units=C * (2 + 3 * components), init_scale=init_scale, vcfg=vcfg)
            assert x.shape == [B, H, W, C * (2 + 3 * components)]
            x = tf.reshape(x, [B, H, W, C, 2 + 3 * components])

            s, t = tf.tanh(x[:, :, :, :, 0]), x[:, :, :, :, 1]
            ml_logits, ml_means, ml_logscales = tf.split(x[:, :, :, :, 2:], 3, axis=4)
            assert s.shape == t.shape == [B, H, W, C]
            assert ml_logits.shape == ml_means.shape == ml_logscales.shape == [B, H, W, C, components]

            return Compose([
                MixLogisticCDF(logits=ml_logits, means=ml_means, logscales=ml_logscales),
                Inverse(Sigmoid()),
                ElemwiseAffine(scales=tf.exp(s), logscales=s, biases=t),
            ])

        self.template = tf.make_template(self.__class__.__name__, f)
//  [python]
def prototypical_layer(nclasses, x_train, y_train, x_test, phi, ext_wts=None):
  """Computes the prototypes, cluster centers.
  Args:
    x_train: [N, ...], Train data.
    y_train: [N], Train class labels.
    x_test: [N, ...], Test data.
    phi: Feature extractor function.
  Returns:
    logits: [N, K], Test prediction.
  """
  protos = [None] * nclasses
  x_all = concat([x_train, x_test], 0)
  h, _ = phi(x_all, reuse=None, is_training=True, ext_wts=ext_wts)
  num_x_train = tf.shape(x_train)[0]
  h_train = h[:num_x_train, :]
  h_test = h[num_x_train:, :]
  for kk in range(nclasses):
    ksel = tf.expand_dims(tf.cast(tf.equal(y_train, kk), h_train.dtype),
                          1)  # [N, 1]
    protos[kk] = tf.reduce_sum(h_train * ksel, [0], keep_dims=True)  # [N, D]
    protos[kk] /= tf.reduce_sum(ksel)
    protos[kk] = tf.Print(protos[kk], [
        'proto', tf.reduce_mean(protos[kk]), tf.reduce_max(protos[kk]),
        tf.reduce_min(protos[kk])
    ])
  protos = concat(protos, 0)  # [K, D]
  logits = compute_logits(protos, h_test)
  return logits
//  [python]
def prototypical_layer(nclasses, x_train, y_train, x_test, phi, ext_wts=None):
  """Computes the prototypes, cluster centers.
  Args:
    x_train: [N, ...], Train data.
    y_train: [N], Train class labels.
    x_test: [N, ...], Test data.
    phi: Feature extractor function.
  Returns:
    logits: [N, K], Test prediction.
  """
  protos = [None] * nclasses
  x_all = concat([x_train, x_test], 0)
  h, _ = phi(x_all, reuse=None, is_training=True, ext_wts=ext_wts)
  num_x_train = tf.shape(x_train)[0]
  h_train = h[:num_x_train, :]
  h_test = h[num_x_train:, :]
  for kk in range(nclasses):
    ksel = tf.expand_dims(tf.cast(tf.equal(y_train, kk), h_train.dtype),
                          1)  # [N, 1]
    protos[kk] = tf.reduce_sum(h_train * ksel, [0], keep_dims=True)  # [N, D]
    protos[kk] /= tf.reduce_sum(ksel)
    protos[kk] = tf.print('proto', tf.reduce_mean(protos[kk]), tf.reduce_max(protos[kk]),
        tf.reduce_min(protos[kk]))
  protos = concat(protos, 0)  # [K, D]
  logits = compute_logits(protos, h_test)
  return logits