jiahuimbzuai/CoQuIR huggingface dataset:
//  [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 _predict_op(self, logits, **kwargs):
        # [batch_size, seq_len, n_groups]
        logits = tf.transpose(logits, perm=[0, 2, 1])
        probs = tf.nn.softmax(logits, axis=-1)
        # [batch_size, seq_len]
        probs = tf.compat.v1.Print(probs, [tf.shape(logits), tf.shape(probs)])
        token_groups = tf.argmax(probs, axis=-1)
        return token_groups, probs
//  [python]
def _predict_op(self, logits, **kwargs):
        # [batch_size, seq_len, n_groups]
        logits = tf.transpose(logits, perm=[0, 2, 1])
        probs = tf.nn.softmax(logits, axis=-1)
        # [batch_size, seq_len]
        probs = tf.print(probs, [tf.shape(logits), tf.shape(probs)])
        token_groups = tf.argmax(probs, axis=-1)
        return token_groups, probs
//  [python]
def add_sigmas_debug(self):

        # Ops to debug variances of random encoders
        enc_sigmas = self.enc_sigmas
        enc_sigmas = tf.Print(
            enc_sigmas,
            [tf.nn.top_k(tf.reshape(enc_sigmas, [-1]), 1).values[0]],
            'Maximal log sigmas:')
        enc_sigmas = tf.Print(
            enc_sigmas,
            [-tf.nn.top_k(tf.reshape(-enc_sigmas, [-1]), 1).values[0]],
            'Minimal log sigmas:')
        self.enc_sigmas = enc_sigmas

        enc_sigmas_t = tf.transpose(self.enc_sigmas)
        max_per_dim = tf.reshape(tf.nn.top_k(enc_sigmas_t, 1).values, [-1, 1])
        min_per_dim = tf.reshape(-tf.nn.top_k(-enc_sigmas_t, 1).values, [-1, 1])
        avg_per_dim = tf.reshape(tf.reduce_mean(enc_sigmas_t, 1), [-1, 1])
        per_dim = tf.concat([min_per_dim, max_per_dim, avg_per_dim], axis=1)
        self.debug_sigmas = per_dim
//  [python]
def add_sigmas_debug(self):

        # Ops to debug variances of random encoders
        enc_sigmas = self.enc_sigmas
        enc_sigmas = tf.print(
            enc_sigmas,
            [tf.nn.top_k(tf.reshape(enc_sigmas, [-1]), 1).values[0]],
            'Maximal log sigmas:')
        enc_sigmas = tf.print(
            enc_sigmas,
            [-tf.nn.top_k(tf.reshape(-enc_sigmas, [-1]), 1).values[0]],
            'Minimal log sigmas:')
        self.enc_sigmas = enc_sigmas

        enc_sigmas_t = tf.transpose(self.enc_sigmas)
        max_per_dim = tf.reshape(tf.nn.top_k(enc_sigmas_t, 1).values, [-1, 1])
        min_per_dim = tf.reshape(-tf.nn.top_k(-enc_sigmas_t, 1).values, [-1, 1])
        avg_per_dim = tf.reshape(tf.reduce_mean(enc_sigmas_t, 1), [-1, 1])
        per_dim = tf.concat([min_per_dim, max_per_dim, avg_per_dim], axis=1)
        self.debug_sigmas = per_dim
//  [python]
def debug_print(self, throughput, to_print=[]):
		return tf.cond(self.train_phase, lambda: throughput, lambda: tf.Print(throughput, to_print + [nop()], "OK!"))
//  [python]
def debug_print(self, throughput, to_print=[]):
		return tf.cond(self.train_phase, lambda: throughput, lambda: tf.print(throughput, *to_print + [nop()], "OK!"))
//  [python]
def Print(self, x, data, message, **kwargs):  # pylint: disable=invalid-name
    """call tf.Print.

    Args:
      x: a LaidOutTensor
      data: a list of LaidOutTensor
      message: a string
      **kwargs: keyword arguments to tf.print
    Returns:
      a LaidOutTensor
    """
    tf.logging.info("PlacementMeshImpl::Print")
    x = x.to_laid_out_tensor()
    new_slices = x.tensor_list[:]
    with tf.device(self._devices[0]):
      new_slices[0] = tf.Print(
          new_slices[0], [t for d in data for t in d.tensor_list],
          message, **kwargs)
    return self.LaidOutTensor(new_slices)
//  [python]
def Print(self, x, data, message, **kwargs):  # pylint: disable=invalid-name
    """call tf.Print.

    Args:
      x: a LaidOutTensor
      data: a list of LaidOutTensor
      message: a string
      **kwargs: keyword arguments to tf.print
    Returns:
      a LaidOutTensor
    """
    tf.logging.info("PlacementMeshImpl::Print")
    x = x.to_laid_out_tensor()
    new_slices = x.tensor_list[:]
    with tf.device(self._devices[0]):
      new_slices[0] = tf.print(
          new_slices[0], [t for d in data for t in d.tensor_list],
          message, **kwargs)
    return self.LaidOutTensor(new_slices)
//  [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.pack(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 compute_loss_from_sptensors(sp_probs, sp_answers):
  """Cross-entropy loss when both answers and probabilities are sparse."""
  # First take log of the probabilities.
  mask = tf.greater(sp_probs.values, 1e-12) # TODO: check this 
  sp_probs = tf.sparse.retain(sp_probs, mask)
  sp_nll = tf.SparseTensor(
      indices=sp_probs.indices,
      values=-tf.log(sp_probs.values),
      dense_shape=sp_probs.dense_shape)
  is_printing=False
  if is_printing:
    tmp_vals = sp_nll.values
    tmp_vals = tf.Print(
        input_=tmp_vals,
        data=[
            tf.shape(sp_probs.indices)[0], sp_probs.indices,
        ],
        message="sp_probs.indices at hop\n",
        first_n=10,
        summarize=50)
    tmp_vals = tf.Print(
        input_=tmp_vals,
        data=[
            tf.shape(sp_answers.indices)[0], sp_answers.indices,
        ],
        message="sp_answers.indices at hop\n",
        first_n=10,
        summarize=50)
    sp_nll = tf.SparseTensor(
        indices=sp_nll.indices,
        values=tmp_vals,
        dense_shape=sp_nll.dense_shape)
  return sp_sp_matmul(sp_nll, sp_answers)
//  [python]
def compute_loss_from_sptensors(sp_probs, sp_answers):
  """Cross-entropy loss when both answers and probabilities are sparse."""
  # First take log of the probabilities.
  mask = tf.greater(sp_probs.values, 1e-12) # TODO: check this 
  sp_probs = tf.sparse.retain(sp_probs, mask)
  sp_nll = tf.SparseTensor(
      indices=sp_probs.indices,
      values=-tf.log(sp_probs.values),
      dense_shape=sp_probs.dense_shape)
  is_printing=False
  if is_printing:
    tmp_vals = sp_nll.values
    tmp_vals = tf.print(
        input_=tmp_vals,
        data=[
            tf.shape(sp_probs.indices)[0], sp_probs.indices,
        ],
        message="sp_probs.indices at hop\n",
        first_n=10,
        summarize=50)
    tmp_vals = tf.print(
        input_=tmp_vals,
        data=[
            tf.shape(sp_answers.indices)[0], sp_answers.indices,
        ],
        message="sp_answers.indices at hop\n",
        first_n=10,
        summarize=50)
    sp_nll = tf.SparseTensor(
        indices=sp_nll.indices,
        values=tmp_vals,
        dense_shape=sp_nll.dense_shape)
  return sp_sp_matmul(sp_nll, sp_answers)
//  [python]
def minibatch(inp, num_kernels=25, kernel_dim=10, scope=None, msg='', reuse_scope=False):
  """
   Borrowed from http://blog.aylien.com/introduction-generative-adversarial-networks-code-tensorflow/
  """
  with tf.variable_scope(scope or 'minibatch_d') as scope:
    scope.set_regularizer(tf.contrib.layers.l2_regularizer(scale=FLAGS.reg_scale))
    if reuse_scope:
      scope.reuse_variables()
  
    inp = tf.Print(inp, [inp],
            '{} inp = '.format(msg), summarize=20, first_n=20)
    x = tf.sigmoid(linear(inp, num_kernels * kernel_dim, scope))
    activation = tf.reshape(x, (-1, num_kernels, kernel_dim))
    activation = tf.Print(activation, [activation],
            '{} activation = '.format(msg), summarize=20, first_n=20)
    diffs = tf.expand_dims(activation, 3) - \
                tf.expand_dims(tf.transpose(activation, [1, 2, 0]), 0)
    diffs = tf.Print(diffs, [diffs],
            '{} diffs = '.format(msg), summarize=20, first_n=20)
    abs_diffs = tf.reduce_sum(tf.abs(diffs), 2)
    abs_diffs = tf.Print(abs_diffs, [abs_diffs],
            '{} abs_diffs = '.format(msg), summarize=20, first_n=20)
    minibatch_features = tf.reduce_sum(tf.exp(-abs_diffs), 2)
    minibatch_features = tf.Print(minibatch_features, [tf.reduce_min(minibatch_features), tf.reduce_max(minibatch_features)],
            '{} minibatch_features (min,max) = '.format(msg), summarize=20, first_n=20)
  return tf.concat( [inp, minibatch_features],1)
//  [python]
def minibatch(inp, num_kernels=25, kernel_dim=10, scope=None, msg='', reuse_scope=False):
  """
   Borrowed from http://blog.aylien.com/introduction-generative-adversarial-networks-code-tensorflow/
  """
  with tf.variable_scope(scope or 'minibatch_d') as scope:
    scope.set_regularizer(tf.contrib.layers.l2_regularizer(scale=FLAGS.reg_scale))
    if reuse_scope:
      scope.reuse_variables()
  
    inp = tf.print(inp,
            '{} inp = '.format(msg), summarize=20, first_n=20)
    x = tf.sigmoid(linear(inp, num_kernels * kernel_dim, scope))
    activation = tf.reshape(x, (-1, num_kernels, kernel_dim))
    activation = tf.print(activation,
            '{} activation = '.format(msg), summarize=20, first_n=20)
    diffs = tf.expand_dims(activation, 3) - \
                tf.expand_dims(tf.transpose(activation, [1, 2, 0]), 0)
    diffs = tf.print(diffs,
            '{} diffs = '.format(msg), summarize=20, first_n=20)
    abs_diffs = tf.reduce_sum(tf.abs(diffs), 2)
    abs_diffs = tf.print(abs_diffs,
            '{} abs_diffs = '.format(msg), summarize=20, first_n=20)
    minibatch_features = tf.reduce_sum(tf.exp(-abs_diffs), 2)
    minibatch_features = tf.print(minibatch_features, [tf.reduce_min(minibatch_features), tf.reduce_max(minibatch_features)],
            '{} minibatch_features (min,max) = '.format(msg), summarize=20, first_n=20)
  return tf.concat( [inp, minibatch_features],1)
//  [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])
            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])
            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 fprop(self, x, **kwargs):
        mean = tf.reduce_mean(x)
        std = tf.sqrt(tf.reduce_mean(tf.square(x - mean)))
        return tf.Print(
            x, [tf.reduce_min(x), mean, tf.reduce_max(x), std], "Print layer"
        )
//  [python]
def fprop(self, x, **kwargs):
        mean = tf.reduce_mean(x)
        std = tf.sqrt(tf.reduce_mean(tf.square(x - mean)))
        return tf.print(
            x, [tf.reduce_min(x), mean, tf.reduce_max(x), std], "Print layer"
        )
//  [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)
            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)
            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 debug_identity(x, name=None):
  if int(os.environ.get('TF_DEBUG', 0)) == 1:
    return tf.Print(x, [
        x.name if name is None else name, tf.reduce_mean(x), tf.reduce_max(x),
        tf.reduce_min(x)
    ])
  else:
    return x
//  [python]
def debug_identity(x, name=None):
  if int(os.environ.get('TF_DEBUG', 0)) == 1:
    return tf.print(x, [
        x.name if name is None else name, tf.reduce_mean(x), tf.reduce_max(x),
        tf.reduce_min(x)
    ])
  else:
    return x
//  [python]
def log_tensor(message, tensor):
    """Log the value of a tensor at graph execution time.

    Warning: this will only work if the tensor is evaluated in your graph.

    Args:
        message: Prefix message string
        tensor: The tensor to evaluate
    """
    tf.Print(tensor, [tensor], message, summarize=10)
//  [python]
def log_tensor(message, tensor):
    """Log the value of a tensor at graph execution time.

    Warning: this will only work if the tensor is evaluated in your graph.

    Args:
        message: Prefix message string
        tensor: The tensor to evaluate
    """
    tf.print(tensor, [tensor], message, summarize=10)
//  [python]
def sp_sp_matmul(sp_a, sp_b):
  """Element-wise multiply two SparseTensors of same shape."""
  is_printing=False
  sp_a = tf.sparse.reorder(sp_a)
  sp_b = tf.sparse.reorder(sp_b)

  # 1. create 1d index maps
  scaling_vector = tf.cumprod(tf.cast(tf.shape(sp_b), sp_b.indices.dtype))
  a1s = tf.linalg.matvec(sp_a.indices, scaling_vector)
  b1s = tf.linalg.matvec(sp_b.indices, scaling_vector)
  
  int_idx_a = tf.setdiff1d(a1s, b1s)[1]
  int_idx_b = tf.setdiff1d(b1s, a1s)[1]

  if is_printing:
    int_idx_a = tf.Print(
        input_=int_idx_a,
        data=[
            tf.shape(int_idx_a)[0],
        ],
        message="int_idx_a.indices at hop\n",
        first_n=10,
        summarize=50)
    int_idx_b = tf.Print(
        input_=int_idx_b,
        data=[
            tf.shape(int_idx_b)[0],
        ],
        message="int_idx_b.indices at hop\n",
        first_n=10,
        summarize=50)


  # 2. get relevant indices of sp_a
  
  to_retain = tf.sparse_to_dense(
      sparse_indices=int_idx_a,
      output_shape=tf.shape(a1s),
      default_value=0.0,
      sparse_values=1.0) < 0.5
  rsp_a = tf.sparse.retain(sp_a, to_retain)
  rsp_a = tf.sparse.reorder(rsp_a)

  # 3. get relevant indices of sp_b
  
  to_retain = tf.sparse_to_dense(
      sparse_indices=int_idx_b,
      output_shape=tf.shape(b1s),
      default_value=0.0,
      sparse_values=1.0) < 0.5
  rsp_b = tf.sparse.retain(sp_b, to_retain)
  rsp_b = tf.sparse.reorder(rsp_b)

  if is_printing:
    tmp_vals = rsp_a.values
    tmp_vals = tf.Print(
        input_= tmp_vals,
        data=[
            tf.shape(rsp_a.indices)[0], rsp_a.indices,
        ],
        message="rsp_a.indices at hop\n",
        first_n=10,
        summarize=100)
    tmp_vals = tf.Print(
        input_= tmp_vals,
        data=[
            tf.shape(rsp_b.indices)[0], rsp_b.indices,
        ],
        message="rsp_b.indices at hop\n",
        first_n=10,
        summarize=100)
    rsp_a = tf.SparseTensor(
          indices=rsp_a.indices,
          values=tmp_vals,
          dense_shape=rsp_a.dense_shape)

  # 4. create output matrix
  output = tf.SparseTensor(
      indices=rsp_b.indices,
      values=rsp_a.values * rsp_b.values,
      dense_shape=sp_b.dense_shape)

  


  return output
//  [python]
def sp_sp_matmul(sp_a, sp_b):
  """Element-wise multiply two SparseTensors of same shape."""
  is_printing=False
  sp_a = tf.sparse.reorder(sp_a)
  sp_b = tf.sparse.reorder(sp_b)

  # 1. create 1d index maps
  scaling_vector = tf.cumprod(tf.cast(tf.shape(sp_b), sp_b.indices.dtype))
  a1s = tf.linalg.matvec(sp_a.indices, scaling_vector)
  b1s = tf.linalg.matvec(sp_b.indices, scaling_vector)
  
  int_idx_a = tf.setdiff1d(a1s, b1s)[1]
  int_idx_b = tf.setdiff1d(b1s, a1s)[1]

  if is_printing:
    int_idx_a = tf.print(
        input_=int_idx_a,
        data=[
            tf.shape(int_idx_a)[0],
        ],
        message="int_idx_a.indices at hop\n",
        first_n=10,
        summarize=50)
    int_idx_b = tf.print(
        input_=int_idx_b,
        data=[
            tf.shape(int_idx_b)[0],
        ],
        message="int_idx_b.indices at hop\n",
        first_n=10,
        summarize=50)


  # 2. get relevant indices of sp_a
  
  to_retain = tf.sparse_to_dense(
      sparse_indices=int_idx_a,
      output_shape=tf.shape(a1s),
      default_value=0.0,
      sparse_values=1.0) < 0.5
  rsp_a = tf.sparse.retain(sp_a, to_retain)
  rsp_a = tf.sparse.reorder(rsp_a)

  # 3. get relevant indices of sp_b
  
  to_retain = tf.sparse_to_dense(
      sparse_indices=int_idx_b,
      output_shape=tf.shape(b1s),
      default_value=0.0,
      sparse_values=1.0) < 0.5
  rsp_b = tf.sparse.retain(sp_b, to_retain)
  rsp_b = tf.sparse.reorder(rsp_b)

  if is_printing:
    tmp_vals = rsp_a.values
    tmp_vals = tf.print(
        input_= tmp_vals,
        data=[
            tf.shape(rsp_a.indices)[0], rsp_a.indices,
        ],
        message="rsp_a.indices at hop\n",
        first_n=10,
        summarize=100)
    tmp_vals = tf.print(
        input_= tmp_vals,
        data=[
            tf.shape(rsp_b.indices)[0], rsp_b.indices,
        ],
        message="rsp_b.indices at hop\n",
        first_n=10,
        summarize=100)
    rsp_a = tf.SparseTensor(
          indices=rsp_a.indices,
          values=tmp_vals,
          dense_shape=rsp_a.dense_shape)

  # 4. create output matrix
  output = tf.SparseTensor(
      indices=rsp_b.indices,
      values=rsp_a.values * rsp_b.values,
      dense_shape=sp_b.dense_shape)

  


  return output
//  [python]
def tf_print(self, var, varStr='null'):
        return tf.Print(var, [var], message=varStr, summarize=100)
//  [python]
def tf_print(self, var, varStr='null'):
        return tf.print(var, [var], message=varStr, summarize=100)
//  [python]
def debugprint(x, name=''):
  """Small wrapper for tf.Print which prints summary statistics."""
  name += '\t' + x.name
  return tf.Print(x,
                  [tf.reduce_min(x), tf.reduce_mean(x), tf.reduce_max(x)],
                  name)
//  [python]
def debugprint(x, name=''):
  """Small wrapper for tf.Print which prints summary statistics."""
  name += '\t' + x.name
  return tf.print(x,
                  [tf.reduce_min(x), tf.reduce_mean(x), tf.reduce_max(x)],
                  name)
//  [python]
def print_stat(x, message=None):
    """ A simple print Op that might be easier to use than :meth:`tf.Print`.
        Use it like: ``x = print_stat(x, message='This is x')``.
    """
    if message is None:
        message = x.op.name
    lst = [tf.shape(x), tf.reduce_mean(x)]
    if x.dtype.is_floating:
        lst.append(rms(x))
    return tf.Print(x, lst + [x], summarize=20,
                    message=message, name='print_' + x.op.name)
//  [python]
def print_stat(x, message=None):
    """ A simple print Op that might be easier to use than :meth:`tf.Print`.
        Use it like: ``x = print_stat(x, message='This is x')``.
    """
    if message is None:
        message = x.op.name
    lst = [tf.shape(x), tf.reduce_mean(x)]
    if x.dtype.is_floating:
        lst.append(rms(x))
    return tf.print(x, lst + [x], summarize=20,
                    message=message, name='print_' + x.op.name)
//  [python]
def expert_dot_product(q, k, v, info_q, info_k):
  """Perform dot product on a subset of the sequence.

  Can add a mask to the attention to prevent sequences to attend to each other
  and to prevent attention to the future.

  Args:
    q (tf.Tensor): Queries of shape [length_expert_q, depth_k]
    k (tf.Tensor): Keys of shape [length_expert_k, depth_k]
    v (tf.Tensor): Values of shape [length_expert_k, depth_v]
    info_q (BatchInfo): Batch info for queries. If None, no mask is added
    info_k (BatchInfo): Batch info for keys

  Returns:
    tf.Tensor: dot product attention output ([length_expert_q, depth_v])
  """

  length_q = common_layers.shape_list(q)[0]
  length_k = common_layers.shape_list(k)[0]
  depth_v = v.get_shape().as_list()[-1]

  # Create the mask
  bias = attention_bias_coordinates(info_q.coordinates, info_k.coordinates)
  if info_k.order is not None:
    bias += attention_bias_future(info_q.order, info_k.order)

  # Restore batch and head dimension
  q, k, v = [tf.expand_dims(tf.expand_dims(t, 0), 0) for t in (q, k, v)]

  def is_zero():
    zeros = tf.zeros(shape=[1, 1, length_q, depth_v], dtype=tf.float32)
    zeros = tf.Print(zeros, [length_k, length_q], "length_k/length_q: ")
    return zeros

  def is_not_zero():
    return dot_product_attention(
        q,
        k,
        v,
        bias=bias,
        # No image summary to avoid "Retval[0] does not have value" (because
        # inside a condition)
        make_image_summary=False,
    )

  # TODO(epot): Should make sure a query gets at least one key. Because the
  # different sequences of a batch are merged, it's possible that a
  # query from a sequence only receive memory from another sequence, so
  # with the mask, the query will perform a softmax on -infinity values.
  # A hack could be to add at least one sequence of each batch on each group so
  # the query can attend to at least one element.
  # Softmax(Q.K)*V
  v_out = tf.cond(
      tf.logical_or(tf.equal(length_q, 0), tf.equal(length_k, 0)),
      is_zero,
      is_not_zero,
  )

  # Remove batch and head dimension
  v_out = tf.squeeze(v_out, axis=0)
  v_out = tf.squeeze(v_out, axis=0)
  return v_out
//  [python]
def expert_dot_product(q, k, v, info_q, info_k):
  """Perform dot product on a subset of the sequence.

  Can add a mask to the attention to prevent sequences to attend to each other
  and to prevent attention to the future.

  Args:
    q (tf.Tensor): Queries of shape [length_expert_q, depth_k]
    k (tf.Tensor): Keys of shape [length_expert_k, depth_k]
    v (tf.Tensor): Values of shape [length_expert_k, depth_v]
    info_q (BatchInfo): Batch info for queries. If None, no mask is added
    info_k (BatchInfo): Batch info for keys

  Returns:
    tf.Tensor: dot product attention output ([length_expert_q, depth_v])
  """

  length_q = common_layers.shape_list(q)[0]
  length_k = common_layers.shape_list(k)[0]
  depth_v = v.get_shape().as_list()[-1]

  # Create the mask
  bias = attention_bias_coordinates(info_q.coordinates, info_k.coordinates)
  if info_k.order is not None:
    bias += attention_bias_future(info_q.order, info_k.order)

  # Restore batch and head dimension
  q, k, v = [tf.expand_dims(tf.expand_dims(t, 0), 0) for t in (q, k, v)]

  def is_zero():
    zeros = tf.zeros(shape=[1, 1, length_q, depth_v], dtype=tf.float32)
    zeros = tf.print(zeros, [length_k, length_q], "length_k/length_q: ")
    return zeros

  def is_not_zero():
    return dot_product_attention(
        q,
        k,
        v,
        bias=bias,
        # No image summary to avoid "Retval[0] does not have value" (because
        # inside a condition)
        make_image_summary=False,
    )

  # TODO(epot): Should make sure a query gets at least one key. Because the
  # different sequences of a batch are merged, it's possible that a
  # query from a sequence only receive memory from another sequence, so
  # with the mask, the query will perform a softmax on -infinity values.
  # A hack could be to add at least one sequence of each batch on each group so
  # the query can attend to at least one element.
  # Softmax(Q.K)*V
  v_out = tf.cond(
      tf.logical_or(tf.equal(length_q, 0), tf.equal(length_k, 0)),
      is_zero,
      is_not_zero,
  )

  # Remove batch and head dimension
  v_out = tf.squeeze(v_out, axis=0)
  v_out = tf.squeeze(v_out, axis=0)
  return v_out
//  [python]
def print_(var, name: str, first_n=10, summarize=5):
	"""Util for debugging, by printing values of tf.Variable `var` during training"""

	# name = (next(k for k, v in globals().items() if v == var) # get name automagically
	# 		if name is None else name) # TODO make work for list ?

	# name = (next(k for k, v in globals().items() if id(v) == id(var))
	# 		if name is None else name)
	# print(name)
	# return ([k for k, v in globals().items() if id(v) == id(var)]
	# 		if name is None else name)

	try:
		return tf.Print(var, [var], '{}: '.format(name), first_n=first_n,
						summarize=summarize)

	except(TypeError): # variables are already in a list
		return tf.Print(var, var, '{}: '.format(name), first_n=first_n,
						summarize=summarize)
//  [python]
def print_(var, name: str, first_n=10, summarize=5):
	"""Util for debugging, by printing values of tf.Variable `var` during training"""

	# name = (next(k for k, v in globals().items() if v == var) # get name automagically
	# 		if name is None else name) # TODO make work for list ?

	# name = (next(k for k, v in globals().items() if id(v) == id(var))
	# 		if name is None else name)
	# print(name)
	# return ([k for k, v in globals().items() if id(v) == id(var)]
	# 		if name is None else name)

	try:
		return tf.print(var, '{}: '.format(name), first_n=first_n,
						summarize=summarize)

	except(TypeError): # variables are already in a list
		return tf.print(var, '{}: '.format(name), first_n=first_n,
						summarize=summarize)
//  [python]
def optimizer_callback(self,parameters):
               
        loss_value, loss_u, loss_f = self.loss(X_u_train, u_train, X_f_train)
        
        u_pred = self.evaluate(X_u_test)
        error_vec = np.linalg.norm((u-u_pred),2)/np.linalg.norm(u,2)
        
        tf.print(loss_value, loss_u, loss_f, error_vec)
//  [python]
def run_bias_check(model, batch, target_type, baseline_type):
  """Run bias check."""
  tf.print(f'Running a bias check comparing {target_type} and {baseline_type}.')
  mu = 0.
  s = 0.
  for step in range(1, int(1e6) + 1):
    diff = run_bias_check_step(
        batch,
        model,
        target_type=target_type,
        baseline_type=baseline_type)
    prev_mu = mu
    mu = mu + (diff - mu) / step
    s = s + (diff - mu) * (diff - prev_mu)

    if step % 100 == 0:
      sigma = tf.math.sqrt(s / step)
      z_score = mu / (sigma / tf.math.sqrt(float(step)))
      tf.print(step, 'z_score: ', z_score, 'sigma: ', sigma)
//  [python]
def custom_loss_abs_genbag(y_true_extended, y_pred):
  """L_1 Generalized Bags Loss."""
  y_combined = tf.concat((y_pred, y_true_extended), axis=1)
  cumu_loss = 0
  loss = 0
  # pylint: disable=cell-var-from-loop
  for partition_num in range(META_BATCH_SIZE):
    y_combined_sliced = tf.gather(
        y_combined,
        tf.where(
            tf.equal(y_combined[:, T5],
                     tf.constant(partition_num, dtype=float)))[:, 0])
    y_combined_sliced_pred = tf.gather(
        y_combined_sliced, tf.range(T2), axis=1)
    y_combined_sliced_actual = tf.gather(
        y_combined_sliced, tf.range(T2, T4), axis=1)
    wts = tf.gather(y_combined_sliced, [T4], axis=1)
    wtd_y_combined_sliced_pred = y_combined_sliced_pred * wts
    wtd_y_combined_sliced_actual = y_combined_sliced_actual * wts
    wtd_sum_y_combined_sliced_pred = tf.reduce_sum(
        wtd_y_combined_sliced_pred, axis=0)
    wtd_sum_y_combined_sliced_actual = tf.reduce_sum(
        wtd_y_combined_sliced_actual, axis=0)
    loss = mae(
        y_true=wtd_sum_y_combined_sliced_actual,
        y_pred=wtd_sum_y_combined_sliced_pred)
    cumu_loss = cumu_loss + loss
    tf.print("Step loss: ", cumu_loss)
  return cumu_loss
//  [python]
def custom_loss_abs_genbag(y_true_extended, y_pred):
  """L_1 Generalized Bags Loss."""
  y_combined = tf.concat((y_pred, y_true_extended), axis=1)
  cumu_loss = 0
  loss = 0
  # pylint: disable=cell-var-from-loop
  for partition_num in range(META_BATCH_SIZE):
    y_combined_sliced = tf.gather(
        y_combined,
        tf.where(
            tf.equal(y_combined[:, T5],
                     tf.constant(partition_num, dtype=float)))[:, 0])
    y_combined_sliced_pred = tf.gather(
        y_combined_sliced, tf.range(T2), axis=1)
    y_combined_sliced_actual = tf.gather(
        y_combined_sliced, tf.range(T2, T4), axis=1)
    wts = tf.gather(y_combined_sliced, [T4], axis=1)
    wtd_y_combined_sliced_pred = y_combined_sliced_pred * wts
    wtd_y_combined_sliced_actual = y_combined_sliced_actual * wts
    wtd_sum_y_combined_sliced_pred = tf.reduce_sum(
        wtd_y_combined_sliced_pred, axis=0)
    wtd_sum_y_combined_sliced_actual = tf.reduce_sum(
        wtd_y_combined_sliced_actual, axis=0)
    loss = mae(
        y_true=wtd_sum_y_combined_sliced_actual,
        y_pred=wtd_sum_y_combined_sliced_pred)
    cumu_loss = cumu_loss + loss
    print()
  tf.print("Step loss: ", cumu_loss)
  return cumu_loss
//  [python]
tf.compat.v1.enable_eager_execution()

input = tf.Variable([0.9, 2.5, 2.3, 1.5, -4.5, 0.45], dtype=tf.float32)
output = tf.round(input)
tf.compat.v1.print(output, summarize=-1)
//  [python]
tf.compat.v1.enable_eager_execution()

input = tf.Variable([0.9, 2.5, 2.3, 1.5, -4.5, 0.45], dtype=tf.float32)
output = tf.round(input)
tf.print(output, summarize=-1)
//  [python]
def hf_compute_loss(self, labels, logits):
        loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(
            from_logits=True, reduction=tf.keras.losses.Reduction.NONE
        )
        # make sure only labels that are not equal to -100
        # are taken into account as loss
        if tf.math.reduce_any(labels == -1):
            tf.compat.v1.logging.warning("Using `-1` to mask the loss for the token is deprecated. Please use `-100` instead.")
            active_loss = tf.reshape(labels, (-1,)) != -1
        else:
            active_loss = tf.reshape(labels, (-1,)) != -100
        reduced_logits = tf.boolean_mask(tf.reshape(logits, (-1, shape_list(logits)[2])), active_loss)
        labels = tf.boolean_mask(tf.reshape(labels, (-1,)), active_loss)

        return loss_fn(labels, reduced_logits)
//  [python]
def hf_compute_loss(self, labels, logits):
        loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(
            from_logits=True, reduction=tf.keras.losses.Reduction.NONE
        )
        # make sure only labels that are not equal to -100
        # are taken into account as loss
        if tf.math.reduce_any(labels == -1):
            tf.print("Using `-1` to mask the loss for the token is deprecated. Please use `-100` instead.")
            active_loss = tf.reshape(labels, (-1,)) != -1
        else:
            active_loss = tf.reshape(labels, (-1,)) != -100
        reduced_logits = tf.boolean_mask(tf.reshape(logits, (-1, shape_list(logits)[2])), active_loss)
        labels = tf.boolean_mask(tf.reshape(labels, (-1,)), active_loss)

        return loss_fn(labels, reduced_logits)
//  [python]
def evaluate(self):
        """Predicts resuts for the test dataset"""

        predictions = []
        LOG.info(f'Predicting segmentation map for test dataset')

        for im in self.test_dataset.as_numpy_iterator():
            DataLoader().validate_schema(im[0])
            break

        for image, mask in self.test_dataset:
            tf.print(image)
            # LOG.info(f'Predicting segmentation map {image}')
            predictions.append(self.model.predict(image))
        return predictions
//  [python]
def _fake_quant_weight(self, weight_tensor, reduce_dims, scope, clamp_value=None):
        if clamp_value is None:
            clamp_value = self._clamp_value

        with tf.variable_scope(scope, reuse=tf.AUTO_REUSE):
            with self._graph.gradient_override_map({'Round': 'Identity'}):
                scales = tf.reduce_max(tf.abs(weight_tensor, name=_MNN_mark_start_name), axis=reduce_dims, keep_dims=True) / clamp_value + self._eps
                weight_scales = tf.get_variable(name='weight_scales', initializer=scales, trainable=True)
                tf.add_to_collection(_MNN_variable_collection_name, weight_scales)
                s_scale = custom_scale_grad(weight_scales, tf.size(weight_tensor), clamp_value)
                # scale_factor = 1. / tf.math.sqrt(tf.cast(tf.size(weight_tensor), tf.float32) * clamp_value + 0.)
                # grad_scale = scale_factor * weight_scales
                # s_scale = grad_scale + tf.stop_gradient(weight_scales - grad_scale)
                quant_w = tf.clip_by_value(tf.round(weight_tensor / s_scale), -clamp_value, clamp_value)
                fake_quant_w = quant_w * s_scale

            if _MNN_QAT_DEBUG:
                print_op = tf.Print(scope, [scope, "weight clamp value:", clamp_value])
                with tf.control_dependencies([print_op]):
                    new_weight_tensor = tf.cond(tf.constant(True), lambda: fake_quant_w, lambda: weight_tensor, name=_MNN_mark_end_name)
            else:
                new_weight_tensor = tf.cond(tf.constant(True), lambda: fake_quant_w, lambda: weight_tensor, name=_MNN_mark_end_name)

        self._weight_scale_original_tensor_map[weight_scales] = weight_tensor

        return new_weight_tensor, weight_scales
//  [python]
def _fake_quant_weight(self, weight_tensor, reduce_dims, scope, clamp_value=None):
        if clamp_value is None:
            clamp_value = self._clamp_value

        with tf.variable_scope(scope, reuse=tf.AUTO_REUSE):
            with self._graph.gradient_override_map({'Round': 'Identity'}):
                scales = tf.reduce_max(tf.abs(weight_tensor, name=_MNN_mark_start_name), axis=reduce_dims, keep_dims=True) / clamp_value + self._eps
                weight_scales = tf.get_variable(name='weight_scales', initializer=scales, trainable=True)
                tf.add_to_collection(_MNN_variable_collection_name, weight_scales)
                s_scale = custom_scale_grad(weight_scales, tf.size(weight_tensor), clamp_value)
                # scale_factor = 1. / tf.math.sqrt(tf.cast(tf.size(weight_tensor), tf.float32) * clamp_value + 0.)
                # grad_scale = scale_factor * weight_scales
                # s_scale = grad_scale + tf.stop_gradient(weight_scales - grad_scale)
                quant_w = tf.clip_by_value(tf.round(weight_tensor / s_scale), -clamp_value, clamp_value)
                fake_quant_w = quant_w * s_scale

            if _MNN_QAT_DEBUG:
                print_op = tf.print(scope, "weight clamp value:", clamp_value)
                with tf.control_dependencies([print_op]):
                    new_weight_tensor = tf.cond(tf.constant(True), lambda: fake_quant_w, lambda: weight_tensor, name=_MNN_mark_end_name)
            else:
                new_weight_tensor = tf.cond(tf.constant(True), lambda: fake_quant_w, lambda: weight_tensor, name=_MNN_mark_end_name)

        self._weight_scale_original_tensor_map[weight_scales] = weight_tensor

        return new_weight_tensor, weight_scales
//  [python]
def on_epoch_end(self, epoch, logs=None):
        print('KGE: epoch=', epoch + 1, ', loss=', logs['loss'], sep='')
//  [python]
def on_epoch_end(self, epoch, logs=None):
        tf.print('KGE: epoch=', epoch + 1, ', loss=', logs['loss'], sep='')
//  [python]
def train_one_step(X, y):    
    with tf.GradientTape() as tape:
        y_pred = model(X)
        loss = tf.keras.losses.sparse_categorical_crossentropy(y_true=y, y_pred=y_pred)
        loss = tf.reduce_mean(loss)
        tf.compat.v1.logging.info("loss %s", loss)
    grads = tape.gradient(loss, model.variables)    
    optimizer.apply_gradients(grads_and_vars=zip(grads, model.variables))
//  [python]
def train_one_step(X, y):    
    with tf.GradientTape() as tape:
        y_pred = model(X)
        loss = tf.keras.losses.sparse_categorical_crossentropy(y_true=y, y_pred=y_pred)
        loss = tf.reduce_mean(loss)
        # 注意这里使用了TensorFlow内置的tf.print()。@tf.function不支持Python内置的print方法
        tf.print("loss", loss)
    grads = tape.gradient(loss, model.variables)    
    optimizer.apply_gradients(grads_and_vars=zip(grads, model.variables))
//  [python]
def reset_model_neurons(self):
    if self.reset_critic:
      tf.print('critic reset')
      self.reset_dead_neurons(
          self._critic_optimizer,
          self.critic1_act,
          self.critic1_act_grads,
          self._critic_network_1,
          self._target_critic_network_1,
      )
      self.reset_dead_neurons(
          self._critic_optimizer,
          self.critic2_act,
          self.critic2_act_grads,
          self._critic_network_2,
          self._target_critic_network_2,
      )
    if self.reset_actor:
      tf.print('actor reset')
      self.reset_dead_neurons(
          self._actor_optimizer,
          self.actor_act,
          self.actor_act_grads,
          self._actor_network,
      )
//  [python]
def _fake_quant_feature(self, feature, scope, op, type, clamp_value=None):
        if clamp_value is None:
            clamp_value = self._clamp_value
        
        with tf.variable_scope(scope, reuse=tf.AUTO_REUSE):
            with self._graph.gradient_override_map({'Round': 'Identity'}):
                feature_abs_max = tf.reduce_max(tf.abs(feature, name=_MNN_mark_start_name))
                if _MNN_QAT_DEBUG:
                    print_op = tf.Print("op:", op.name, type, "feature_abs_max, clamp_value:", [feature_abs_max, clamp_value])
                    with tf.control_dependencies([print_op]):
                        instance_scale = feature_abs_max / clamp_value
                else:
                    instance_scale = feature_abs_max / clamp_value
                moving_average_scale = tf.get_variable(name='moving_average_scale', initializer=0., trainable=True)
                tf.add_to_collection(_MNN_variable_collection_name, moving_average_scale)

                def init():
                    init_assign = tf.assign(moving_average_scale, instance_scale)
                    initialize_op = tf.assign(self._op_scale_initialized[op][type], 1)
                    if _MNN_QAT_DEBUG:
                        print_op = tf.Print("init", moving_average_scale.name, init_assign, [moving_average_scale.name, init_assign])
                        with tf.control_dependencies([print_op, initialize_op]):
                            return tf.identity(init_assign)
                    else:
                        with tf.control_dependencies([initialize_op]):
                            return tf.identity(init_assign)

                def do_nothing():
                    return moving_average_scale

                init_moving_scale = tf.cond(tf.equal(self._op_scale_initialized[op][type], 0), init, do_nothing)
                self._feature_scale_init_op[moving_average_scale] = init_moving_scale

                def fake_quant():
                    with tf.control_dependencies([init_moving_scale]):
                        s_scale = custom_scale_grad(moving_average_scale, tf.size(feature), clamp_value)
                        quant_x = tf.clip_by_value(tf.round(feature / s_scale), -clamp_value, clamp_value)
                        fake_quant_x = quant_x * s_scale
                        return fake_quant_x

                new_feature = tf.cond(tf.constant(True), fake_quant, lambda:feature, name=_MNN_mark_end_name)
        
        self._feature_scale_original_tensor_dims_map[moving_average_scale] = [feature, feature.shape.as_list()[-1]]
        
        return new_feature, moving_average_scale, feature.shape.as_list()[-1], init_moving_scale