jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def rouge_2_fscore(predictions, labels, **unused_kwargs):
  """ROUGE-2 F1 score computation between labels and predictions.
  This is an approximate ROUGE scoring method since we do not glue word pieces
  or decode the ids and tokenize the output.
  Args:
    predictions: tensor, model predictions
    labels: tensor, gold output.
  Returns:
    rouge2_fscore: approx rouge-2 f1 score.
  """

  outputs = tf.to_int32(tf.argmax(predictions, axis=-1))
  # Convert the outputs and labels to a [batch_size, input_length] tensor.
  outputs = tf.squeeze(outputs, axis=[-1, -2])
  labels = tf.squeeze(labels, axis=[-1, -2])
  rouge_2_f_score = tf.py_func(rouge_n, (outputs, labels), tf.float32)
  return rouge_2_f_score, tf.constant(1.0)
//  [python]
def rouge_2_fscore(predictions, labels, **unused_kwargs):
  """ROUGE-2 F1 score computation between labels and predictions.
  This is an approximate ROUGE scoring method since we do not glue word pieces
  or decode the ids and tokenize the output.
  Args:
    predictions: tensor, model predictions
    labels: tensor, gold output.
  Returns:
    rouge2_fscore: approx rouge-2 f1 score.
  """

  outputs = tf.cast(tf.argmax(predictions, axis=-1), tf.int32)
  # Convert the outputs and labels to a [batch_size, input_length] tensor.
  outputs = tf.squeeze(outputs, axis=[-1, -2])
  labels = tf.squeeze(labels, axis=[-1, -2])
  rouge_2_f_score = tf.py_func(rouge_n, (outputs, labels), tf.float32)
  return rouge_2_f_score, tf.constant(1.0)
//  [python]
def unit(w, sparsity):
  """Unit-level magnitude pruning."""
  w_shape = common_layers.shape_list(w)
  count = tf.to_int32(w_shape[-1] * sparsity)
  mask = common_layers.unit_targeting(w, count)
  return (1 - mask) * w
//  [python]
def unit(w, sparsity):
  """Unit-level magnitude pruning."""
  w_shape = common_layers.shape_list(w)
  count = tf.cast(w_shape[-1] * sparsity, tf.int32)
  mask = common_layers.unit_targeting(w, count)
  return (1 - mask) * w
//  [python]
def get_data_and_params():
  """Set up input dataset and variables."""
  (train_x, train_y), _ = tf.keras.datasets.mnist.load_data()
  tf.set_random_seed(0)
  hparams = contrib_training.HParams(
      batch_size=200,
      learning_rate=0.1,
      train_steps=101,
  )
  dataset = tf.data.Dataset.from_tensor_slices((train_x, train_y))
  dataset = dataset.repeat()
  dataset = dataset.shuffle(hparams.batch_size * 10)
  dataset = dataset.batch(hparams.batch_size)

  def reshape_ex(x, y):
    return (tf.to_float(tf.reshape(x, (-1, 28 * 28))) / 256.0,
            tf.one_hot(tf.squeeze(y), 10))

  dataset = dataset.map(reshape_ex)
  w = tf.get_variable('w0', (28 * 28, 10))
  b = tf.get_variable('b0', (10,), initializer=tf.zeros_initializer())
  opt = tf.train.GradientDescentOptimizer(hparams.learning_rate)
  return dataset, opt, hparams, w, b
//  [python]
def get_data_and_params():
  """Set up input dataset and variables."""
  (train_x, train_y), _ = tf.keras.datasets.mnist.load_data()
  tf.set_random_seed(0)
  hparams = contrib_training.HParams(
      batch_size=200,
      learning_rate=0.1,
      train_steps=101,
  )
  dataset = tf.data.Dataset.from_tensor_slices((train_x, train_y))
  dataset = dataset.repeat()
  dataset = dataset.shuffle(hparams.batch_size * 10)
  dataset = dataset.batch(hparams.batch_size)

  def reshape_ex(x, y):
    return (tf.cast(tf.reshape(x, (-1, 28 * 28)), tf.float32) / 256.0,
            tf.one_hot(tf.squeeze(y), 10))

  dataset = dataset.map(reshape_ex)
  w = tf.get_variable('w0', (28 * 28, 10))
  b = tf.get_variable('b0', (10,), initializer=tf.zeros_initializer())
  opt = tf.train.GradientDescentOptimizer(hparams.learning_rate)
  return dataset, opt, hparams, w, b
//  [python]
def cls_ohem(cls_prob, label):
    zeros = tf.zeros_like(label)
    # label=-1 --> label=0net_factory

    # pos -> 1, neg -> 0, others -> 0
    label_filter_invalid = tf.where(tf.less(label, 0), zeros, label)
    num_cls_prob = tf.size(cls_prob)
    cls_prob_reshape = tf.reshape(cls_prob, [num_cls_prob, -1])
    label_int = tf.cast(label_filter_invalid, tf.int32)
    # get the number of rows of class_prob
    num_row = tf.to_int32(cls_prob.get_shape()[0])
    # row = [0,2,4.....]
    row = tf.range(num_row) * 2
    indices_ = row + label_int
    label_prob = tf.squeeze(tf.gather(cls_prob_reshape, indices_))
    loss = -tf.log(label_prob + 1e-10)
    zeros = tf.zeros_like(label_prob, dtype=tf.float32)
    ones = tf.ones_like(label_prob, dtype=tf.float32)
    # set pos and neg to be 1, rest to be 0
    valid_inds = tf.where(label < zeros, zeros, ones)
    # get the number of POS and NEG examples
    num_valid = tf.reduce_sum(valid_inds)

    keep_num = tf.cast(num_valid * num_keep_radio, dtype=tf.int32)
    # FILTER OUT PART AND LANDMARK DATA
    loss = loss * valid_inds
    loss, _ = tf.nn.top_k(loss, k=keep_num)
    return tf.reduce_mean(loss)
//  [python]
def cls_ohem(cls_prob, label):
    zeros = tf.zeros_like(label)
    # label=-1 --> label=0net_factory

    # pos -> 1, neg -> 0, others -> 0
    label_filter_invalid = tf.where(tf.less(label, 0), zeros, label)
    num_cls_prob = tf.size(cls_prob)
    cls_prob_reshape = tf.reshape(cls_prob, [num_cls_prob, -1])
    label_int = tf.cast(label_filter_invalid, tf.int32)
    # get the number of rows of class_prob
    num_row = tf.cast(cls_prob.get_shape()[0], tf.int32)
    # row = [0,2,4.....]
    row = tf.range(num_row) * 2
    indices_ = row + label_int
    label_prob = tf.squeeze(tf.gather(cls_prob_reshape, indices_))
    loss = -tf.log(label_prob + 1e-10)
    zeros = tf.zeros_like(label_prob, dtype=tf.float32)
    ones = tf.ones_like(label_prob, dtype=tf.float32)
    # set pos and neg to be 1, rest to be 0
    valid_inds = tf.where(label < zeros, zeros, ones)
    # get the number of POS and NEG examples
    num_valid = tf.reduce_sum(valid_inds)

    keep_num = tf.cast(num_valid * num_keep_radio, dtype=tf.int32)
    # FILTER OUT PART AND LANDMARK DATA
    loss = loss * valid_inds
    loss, _ = tf.nn.top_k(loss, k=keep_num)
    return tf.reduce_mean(loss)
//  [python]
def get_attention_bias(sequence_length, maxlen=None):
  """Create attention bias so attention is not applied at padding position."""
  # attention_bias: [batch, 1, 1, memory_length]
  mask = tf.sequence_mask(sequence_length, maxlen=maxlen)
  nonpadding = tf.to_float(mask)
  invert_sequence_mask = tf.to_float(tf.logical_not(mask))
  attention_bias = common_attention.attention_bias_ignore_padding(
      invert_sequence_mask)
  return nonpadding, attention_bias
//  [python]
def get_attention_bias(sequence_length, maxlen=None):
  """Create attention bias so attention is not applied at padding position."""
  # attention_bias: [batch, 1, 1, memory_length]
  mask = tf.sequence_mask(sequence_length, maxlen=maxlen)
  nonpadding = tf.cast(mask, tf.float32)
  invert_sequence_mask = tf.cast(tf.logical_not(mask), tf.float32)
  attention_bias = common_attention.attention_bias_ignore_padding(
      invert_sequence_mask)
  return nonpadding, attention_bias
//  [python]
def positional_encoding(inputs,
                        maxlen,
                        masking=True,
                        scope="positional_encoding"):
    '''Sinusoidal Positional_Encoding. See 3.5
    inputs: 3d tensor. (N, T, E)
    maxlen: scalar. Must be >= T
    masking: Boolean. If True, padding positions are set to zeros.
    scope: Optional scope for `variable_scope`.

    returns
    3d tensor that has the same shape as inputs.
    '''

    E = inputs.get_shape().as_list()[-1]  # static
    N, T = tf.shape(inputs)[0], tf.shape(inputs)[1]  # dynamic
    with tf.variable_scope(scope, reuse=tf.AUTO_REUSE):
        # position indices
        position_ind = tf.tile(tf.expand_dims(
            tf.range(T), 0), [N, 1])  # (N, T)

        # First part of the PE function: sin and cos argument
        position_enc = np.array([
            [pos / np.power(10000, (i-i % 2)/E) for i in range(E)]
            for pos in range(maxlen)])

        # Second part, apply the cosine to even columns and sin to odds.
        position_enc[:, 0::2] = np.sin(position_enc[:, 0::2])  # dim 2i
        position_enc[:, 1::2] = np.cos(position_enc[:, 1::2])  # dim 2i+1
        position_enc = tf.convert_to_tensor(
            position_enc, tf.float32)  # (maxlen, E)

        # lookup
        outputs = tf.nn.embedding_lookup(position_enc, position_ind)

        # masks
        if masking:
            outputs = tf.where(tf.equal(inputs, 0), inputs, outputs)

        return tf.to_float(outputs)
//  [python]
def positional_encoding(inputs,
                        maxlen,
                        masking=True,
                        scope="positional_encoding"):
    '''Sinusoidal Positional_Encoding. See 3.5
    inputs: 3d tensor. (N, T, E)
    maxlen: scalar. Must be >= T
    masking: Boolean. If True, padding positions are set to zeros.
    scope: Optional scope for `variable_scope`.

    returns
    3d tensor that has the same shape as inputs.
    '''

    E = inputs.get_shape().as_list()[-1]  # static
    N, T = tf.shape(inputs)[0], tf.shape(inputs)[1]  # dynamic
    with tf.variable_scope(scope, reuse=tf.AUTO_REUSE):
        # position indices
        position_ind = tf.tile(tf.expand_dims(
            tf.range(T), 0), [N, 1])  # (N, T)

        # First part of the PE function: sin and cos argument
        position_enc = np.array([
            [pos / np.power(10000, (i-i % 2)/E) for i in range(E)]
            for pos in range(maxlen)])

        # Second part, apply the cosine to even columns and sin to odds.
        position_enc[:, 0::2] = np.sin(position_enc[:, 0::2])  # dim 2i
        position_enc[:, 1::2] = np.cos(position_enc[:, 1::2])  # dim 2i+1
        position_enc = tf.convert_to_tensor(
            position_enc, tf.float32)  # (maxlen, E)

        # lookup
        outputs = tf.nn.embedding_lookup(position_enc, position_ind)

        # masks
        if masking:
            outputs = tf.where(tf.equal(inputs, 0), inputs, outputs)

        return tf.cast(outputs, tf.float32)
//  [python]
def loss(self, logits, labels, regularization, data ,extra_loss=False, regression=False):
        """Adds to the inference model the layers required to generate loss."""
        with tf.name_scope('loss'):
            if regression:
                with tf.name_scope('MSE'):
                    predictions = tf.squeeze(logits)
#                     if self.M:
#                         labels = tf.expand_dims(labels, axis=-1)
                    if hasattr(self, 'train_mask'):
                        predictions = predictions * data[..., -2]
                        labels = labels * data[..., -1]
                    mse = tf.losses.mean_squared_error(labels, predictions)
                    loss = mse
            else:
                with tf.name_scope('cross_entropy'):
                    labels = tf.to_int64(labels)
                    labels_onehot = tf.one_hot(labels, 3)
#                     weights = tf.constant([[0.00102182, 0.95426438, 0.04471379]])
                    if self.weighted:
                        weights = tf.constant([[0.34130685, 318.47388343,  14.93759951]])
                        batch_weights = tf.reshape(tf.matmul(tf.reshape(labels_onehot, [-1,3]), tf.transpose(weights)), 
                                                   [self.batch_size, self.L[0].shape[0]])
#                     batch_weights = tf.reduce_sum(class_weights * onehot_labels, axis=1)
                    cross_entropy = tf.nn.sparse_softmax_cross_entropy_with_logits(logits=logits, labels=labels)
                    if self.weighted:
                        cross_entropy = tf.multiply(batch_weights, cross_entropy)  
#                     cross_entropy = tf.reduce_sum(cross_entropy*batch_weights) / self.batch_size
                    cross_entropy = tf.reduce_mean(cross_entropy)
                    loss = cross_entropy
            with tf.name_scope('regularization'):
                n_weights = np.sum(self.regularizers_size)
                regularization *= tf.add_n(self.regularizers) / n_weights
            loss = loss + regularization
            if extra_loss:
                loss += self.triplet_loss(self.op_descriptor, labels)

            # Summaries for TensorBoard.
            if regression:
                tf.summary.scalar('loss/mse', mse)
            else:
                tf.summary.scalar('loss/cross_entropy', cross_entropy)
            tf.summary.scalar('loss/regularization', regularization)
            tf.summary.scalar('loss/total', loss)
            return loss
//  [python]
def loss(self, logits, labels, regularization, data ,extra_loss=False, regression=False):
        """Adds to the inference model the layers required to generate loss."""
        with tf.name_scope('loss'):
            if regression:
                with tf.name_scope('MSE'):
                    predictions = tf.squeeze(logits)
#                     if self.M:
#                         labels = tf.expand_dims(labels, axis=-1)
                    if hasattr(self, 'train_mask'):
                        predictions = predictions * data[..., -2]
                        labels = labels * data[..., -1]
                    mse = tf.losses.mean_squared_error(labels, predictions)
                    loss = mse
            else:
                with tf.name_scope('cross_entropy'):
                    labels = tf.cast(labels, tf.int64)
                    labels_onehot = tf.one_hot(labels, 3)
#                     weights = tf.constant([[0.00102182, 0.95426438, 0.04471379]])
                    if self.weighted:
                        weights = tf.constant([[0.34130685, 318.47388343,  14.93759951]])
                        batch_weights = tf.reshape(tf.matmul(tf.reshape(labels_onehot, [-1,3]), tf.transpose(weights)), 
                                                   [self.batch_size, self.L[0].shape[0]])
#                     batch_weights = tf.reduce_sum(class_weights * onehot_labels, axis=1)
                    cross_entropy = tf.nn.sparse_softmax_cross_entropy_with_logits(logits=logits, labels=labels)
                    if self.weighted:
                        cross_entropy = tf.multiply(batch_weights, cross_entropy)  
#                     cross_entropy = tf.reduce_sum(cross_entropy*batch_weights) / self.batch_size
                    cross_entropy = tf.reduce_mean(cross_entropy)
                    loss = cross_entropy
            with tf.name_scope('regularization'):
                n_weights = np.sum(self.regularizers_size)
                regularization *= tf.add_n(self.regularizers) / n_weights
            loss = loss + regularization
            if extra_loss:
                loss += self.triplet_loss(self.op_descriptor, labels)

            # Summaries for TensorBoard.
            if regression:
                tf.summary.scalar('loss/mse', mse)
            else:
                tf.summary.scalar('loss/cross_entropy', cross_entropy)
            tf.summary.scalar('loss/regularization', regularization)
            tf.summary.scalar('loss/total', loss)
            return loss
//  [python]
def _compute_word_overlap(context_ids, context_len, question_ids, question_len,
                          reduce_type, weighted, vocab_df):
  """Compute word overlap between question and context ids.

  Args:
    context_ids: <int32> [batch_size, num_contexts, max_context_len]
    context_len: <int32> [batch_size, num_contexts]
    question_ids: <int32> [batch_size, max_question_len]
    question_len: <int32> [batch_size]
    reduce_type: String for reduce type when computing overlap. Choices are: max
      - Allows at most one match per question word. sum - Sums over all matches
      for each question word.
    weighted: Boolean indicate whether or not weight the overlap by IDF.
    vocab_df: Tensor of shape [vocab_size] for word frequency. Computes this at
      the document-level if not given.

  Returns:
    overlap: <float32> [batch_size, num_contexts]

  Raises:
    Exception: If invalid reduce_type is provided.
  """
  # <float> [batch_size, num_contexts, question_len, context_len]
  overlap = tf.to_float(
      _word_overlap_helper(question_ids=question_ids, context_ids=context_ids))

  # <float> [batch_size, question_len]
  question_mask = tf.sequence_mask(
      question_len, tf.shape(question_ids)[1], dtype=tf.float32)

  # <float> [batch_size, num_contexts, context_len]
  context_mask = tf.sequence_mask(
      context_len, tf.shape(context_ids)[2], dtype=tf.float32)

  overlap *= tf.expand_dims(tf.expand_dims(question_mask, 1), -1)
  overlap *= tf.expand_dims(context_mask, 2)

  if weighted:
    if vocab_df is None:
      # Use document-level IDF computed with respect to the current batch.
      flat_context_ids = tf.to_int32(tf.reshape(context_ids, [-1]))

      # <float> [number of unique words]
      vocab_df = tf.bincount(
          flat_context_ids,
          minlength=tf.reduce_max(question_ids) + 1,
          dtype=tf.float32)

      # Replace all zeros with ones.
      vocab_df = tf.where(
          tf.equal(vocab_df, 0), x=tf.ones_like(vocab_df), y=vocab_df)

    # <float>[batch_size, question_len] expanded to
    # <float> [batch_size, 1, question_len, 1]
    question_df = tf.gather(vocab_df, question_ids)
    question_df = tf.expand_dims(tf.expand_dims(question_df, 1), -1)

    # <float> [batch_size, num_contexts, question_len, context_len]
    overlap = tf.divide(tf.to_float(overlap), question_df)

  if reduce_type == "max":
    # <float> [batch_size, num_contexts]
    overlap = tf.reduce_sum(tf.reduce_max(overlap, axis=[3]), axis=[2])
  elif reduce_type == "sum":
    # <float> [batch_size, num_contexts]
    overlap = tf.reduce_sum(overlap, axis=[2, 3])
  else:
    raise Exception("Reduce type %s is invalid." % reduce_type)

  return overlap
//  [python]
def _compute_word_overlap(context_ids, context_len, question_ids, question_len,
                          reduce_type, weighted, vocab_df):
  """Compute word overlap between question and context ids.

  Args:
    context_ids: <int32> [batch_size, num_contexts, max_context_len]
    context_len: <int32> [batch_size, num_contexts]
    question_ids: <int32> [batch_size, max_question_len]
    question_len: <int32> [batch_size]
    reduce_type: String for reduce type when computing overlap. Choices are: max
      - Allows at most one match per question word. sum - Sums over all matches
      for each question word.
    weighted: Boolean indicate whether or not weight the overlap by IDF.
    vocab_df: Tensor of shape [vocab_size] for word frequency. Computes this at
      the document-level if not given.

  Returns:
    overlap: <float32> [batch_size, num_contexts]

  Raises:
    Exception: If invalid reduce_type is provided.
  """
  # <float> [batch_size, num_contexts, question_len, context_len]
  overlap = tf.cast(
      _word_overlap_helper(question_ids=question_ids, context_ids=context_ids), tf.float32)

  # <float> [batch_size, question_len]
  question_mask = tf.sequence_mask(
      question_len, tf.shape(question_ids)[1], dtype=tf.float32)

  # <float> [batch_size, num_contexts, context_len]
  context_mask = tf.sequence_mask(
      context_len, tf.shape(context_ids)[2], dtype=tf.float32)

  overlap *= tf.expand_dims(tf.expand_dims(question_mask, 1), -1)
  overlap *= tf.expand_dims(context_mask, 2)

  if weighted:
    if vocab_df is None:
      # Use document-level IDF computed with respect to the current batch.
      flat_context_ids = tf.to_int32(tf.reshape(context_ids, [-1]))

      # <float> [number of unique words]
      vocab_df = tf.bincount(
          flat_context_ids,
          minlength=tf.reduce_max(question_ids) + 1,
          dtype=tf.float32)

      # Replace all zeros with ones.
      vocab_df = tf.where(
          tf.equal(vocab_df, 0), x=tf.ones_like(vocab_df), y=vocab_df)

    # <float>[batch_size, question_len] expanded to
    # <float> [batch_size, 1, question_len, 1]
    question_df = tf.gather(vocab_df, question_ids)
    question_df = tf.expand_dims(tf.expand_dims(question_df, 1), -1)

    # <float> [batch_size, num_contexts, question_len, context_len]
    overlap = tf.divide(tf.cast(overlap, tf.float32), question_df)

  if reduce_type == "max":
    # <float> [batch_size, num_contexts]
    overlap = tf.reduce_sum(tf.reduce_max(overlap, axis=[3]), axis=[2])
  elif reduce_type == "sum":
    # <float> [batch_size, num_contexts]
    overlap = tf.reduce_sum(overlap, axis=[2, 3])
  else:
    raise Exception("Reduce type %s is invalid." % reduce_type)

  return overlap
//  [python]
def _build_interest_simi(self):
    user_emb_num = self._prediction_dict['user_emb_num']
    high_capsule_mask = tf.sequence_mask(
        user_emb_num, self._model_config.capsule_config.max_k)

    user_interests = self._prediction_dict['user_interests']
    high_capsule_mask = tf.to_float(high_capsule_mask[:, :, None])
    user_interests = self.norm(user_interests) * high_capsule_mask

    user_feature_sum_sqr = tf.square(tf.reduce_sum(user_interests, axis=1))
    user_feature_sqr_sum = tf.reduce_sum(tf.square(user_interests), axis=1)
    interest_simi = user_feature_sum_sqr - user_feature_sqr_sum

    high_capsules = self._prediction_dict['high_capsules']
    high_capsules = self.norm(high_capsules) * high_capsule_mask
    high_capsule_sum_sqr = tf.square(tf.reduce_sum(high_capsules, axis=1))
    high_capsule_sqr_sum = tf.reduce_sum(tf.square(high_capsules), axis=1)
    high_capsule_simi = high_capsule_sum_sqr - high_capsule_sqr_sum

    # normalize by interest number
    interest_div = tf.maximum(
        tf.to_float(user_emb_num * (user_emb_num - 1)), 1.0)
    interest_simi = tf.reduce_sum(interest_simi, axis=1) / interest_div

    high_capsule_simi = tf.reduce_sum(high_capsule_simi, axis=1) / interest_div

    # normalize by batch_size
    multi_interest = tf.to_float(user_emb_num > 1)
    sum_interest_simi = tf.reduce_sum(
        (interest_simi + 1) * multi_interest) / 2.0
    sum_div = tf.maximum(tf.reduce_sum(multi_interest), 1.0)
    avg_interest_simi = sum_interest_simi / sum_div

    sum_capsule_simi = tf.reduce_sum(
        (high_capsule_simi + 1) * multi_interest) / 2.0
    avg_capsule_simi = sum_capsule_simi / sum_div

    tf.summary.scalar('interest_similarity', avg_interest_simi)
    tf.summary.scalar('capsule_similarity', avg_capsule_simi)
    return avg_interest_simi, avg_capsule_simi
//  [python]
def _build_interest_simi(self):
    user_emb_num = self._prediction_dict['user_emb_num']
    high_capsule_mask = tf.sequence_mask(
        user_emb_num, self._model_config.capsule_config.max_k)

    user_interests = self._prediction_dict['user_interests']
    high_capsule_mask = tf.cast(high_capsule_mask[:, :, None], tf.float32)
    user_interests = self.norm(user_interests) * high_capsule_mask

    user_feature_sum_sqr = tf.square(tf.reduce_sum(user_interests, axis=1))
    user_feature_sqr_sum = tf.reduce_sum(tf.square(user_interests), axis=1)
    interest_simi = user_feature_sum_sqr - user_feature_sqr_sum

    high_capsules = self._prediction_dict['high_capsules']
    high_capsules = self.norm(high_capsules) * high_capsule_mask
    high_capsule_sum_sqr = tf.square(tf.reduce_sum(high_capsules, axis=1))
    high_capsule_sqr_sum = tf.reduce_sum(tf.square(high_capsules), axis=1)
    high_capsule_simi = high_capsule_sum_sqr - high_capsule_sqr_sum

    # normalize by interest number
    interest_div = tf.maximum(
        tf.cast(user_emb_num * (user_emb_num - 1), tf.float32), 1.0)
    interest_simi = tf.reduce_sum(interest_simi, axis=1) / interest_div

    high_capsule_simi = tf.reduce_sum(high_capsule_simi, axis=1) / interest_div

    # normalize by batch_size
    multi_interest = tf.cast(user_emb_num > 1, tf.float32)
    sum_interest_simi = tf.reduce_sum(
        (interest_simi + 1) * multi_interest) / 2.0
    sum_div = tf.maximum(tf.reduce_sum(multi_interest), 1.0)
    avg_interest_simi = sum_interest_simi / sum_div

    sum_capsule_simi = tf.reduce_sum(
        (high_capsule_simi + 1) * multi_interest) / 2.0
    avg_capsule_simi = sum_capsule_simi / sum_div

    tf.summary.scalar('interest_similarity', avg_interest_simi)
    tf.summary.scalar('capsule_similarity', avg_capsule_simi)
    return avg_interest_simi, avg_capsule_simi
//  [python]
def _smooth_l1_loss_base(self, bbox_pred, bbox_targets, sigma=1.0):
        '''
        :param bbox_pred: [-1, 4] in RPN. [-1, cls_num+1, 4] in Fast-rcnn
        :param bbox_targets: shape is same as bbox_pred
        :param sigma:
        :return:
        '''
        sigma_2 = sigma ** 2

        box_diff = bbox_pred - bbox_targets

        abs_box_diff = tf.abs(box_diff)

        smoothL1_sign = tf.stop_gradient(
            tf.to_float(tf.less(abs_box_diff, 1. / sigma_2)))
        loss_box = tf.pow(box_diff, 2) * (sigma_2 / 2.0) * smoothL1_sign \
                   + (abs_box_diff - (0.5 / sigma_2)) * (1.0 - smoothL1_sign)
        return loss_box
//  [python]
def _smooth_l1_loss_base(self, bbox_pred, bbox_targets, sigma=1.0):
        '''
        :param bbox_pred: [-1, 4] in RPN. [-1, cls_num+1, 4] in Fast-rcnn
        :param bbox_targets: shape is same as bbox_pred
        :param sigma:
        :return:
        '''
        sigma_2 = sigma ** 2

        box_diff = bbox_pred - bbox_targets

        abs_box_diff = tf.abs(box_diff)

        smoothL1_sign = tf.stop_gradient(
            tf.cast(tf.less(abs_box_diff, 1. / sigma_2), tf.float32))
        loss_box = tf.pow(box_diff, 2) * (sigma_2 / 2.0) * smoothL1_sign \
                   + (abs_box_diff - (0.5 / sigma_2)) * (1.0 - smoothL1_sign)
        return loss_box
//  [python]
def build_train_op(loss, is_tpu, opt_cfg, trainable_vars=None):
  """Build the optimizer, compute gradients, and build training op."""
  trainable_vars = trainable_vars or tf.trainable_variables()

  def decay_fn(learning_rate, global_step):
    # Linear warmup from 0.

    lr_decay = tf.minimum(1.0, tf.to_float(global_step) / opt_cfg.warmup_steps)
    return lr_decay * learning_rate

  def optimizer(lr):
    opt = tf.train.AdamOptimizer(
        learning_rate=lr,
        beta1=opt_cfg.beta1,
        beta2=opt_cfg.beta2,
        epsilon=opt_cfg.epsilon)
    if is_tpu:
      opt = contrib_tpu.CrossShardOptimizer(opt)
    return opt

  if opt_cfg.l2_reg:
    tf.logging.info('Applying l2 regularization of %s', opt_cfg.l2_reg)
    decay_costs = []
    for var in trainable_vars:
      decay_costs.append(tf.nn.l2_loss(var))

    loss += tf.multiply(opt_cfg.l2_reg, tf.add_n(decay_costs))

  train_op = contrib_layers.optimize_loss(
      loss=loss,
      global_step=tf.train.get_global_step(),
      learning_rate=opt_cfg.lr,
      learning_rate_decay_fn=decay_fn,
      clip_gradients=opt_cfg.grad_clip_norm,
      optimizer=optimizer,
      variables=trainable_vars,
      gradient_noise_scale=None,
      summaries=[] if is_tpu else [
          'learning_rate',
          'loss',
      ])

  if opt_cfg.ema_decay < 1.0:
    ema = tf.train.ExponentialMovingAverage(decay=opt_cfg.ema_decay)
    maintain_average_op = ema.apply(trainable_vars)
    with tf.control_dependencies([train_op]):
      train_op = tf.group(maintain_average_op)

  return train_op
//  [python]
def build_train_op(loss, is_tpu, opt_cfg, trainable_vars=None):
  """Build the optimizer, compute gradients, and build training op."""
  trainable_vars = trainable_vars or tf.trainable_variables()

  def decay_fn(learning_rate, global_step):
    # Linear warmup from 0.

    lr_decay = tf.minimum(1.0, tf.cast(global_step, tf.float32) / opt_cfg.warmup_steps)
    return lr_decay * learning_rate

  def optimizer(lr):
    opt = tf.train.AdamOptimizer(
        learning_rate=lr,
        beta1=opt_cfg.beta1,
        beta2=opt_cfg.beta2,
        epsilon=opt_cfg.epsilon)
    if is_tpu:
      opt = contrib_tpu.CrossShardOptimizer(opt)
    return opt

  if opt_cfg.l2_reg:
    tf.logging.info('Applying l2 regularization of %s', opt_cfg.l2_reg)
    decay_costs = []
    for var in trainable_vars:
      decay_costs.append(tf.nn.l2_loss(var))

    loss += tf.multiply(opt_cfg.l2_reg, tf.add_n(decay_costs))

  train_op = contrib_layers.optimize_loss(
      loss=loss,
      global_step=tf.train.get_global_step(),
      learning_rate=opt_cfg.lr,
      learning_rate_decay_fn=decay_fn,
      clip_gradients=opt_cfg.grad_clip_norm,
      optimizer=optimizer,
      variables=trainable_vars,
      gradient_noise_scale=None,
      summaries=[] if is_tpu else [
          'learning_rate',
          'loss',
      ])

  if opt_cfg.ema_decay < 1.0:
    ema = tf.train.ExponentialMovingAverage(decay=opt_cfg.ema_decay)
    maintain_average_op = ema.apply(trainable_vars)
    with tf.control_dependencies([train_op]):
      train_op = tf.group(maintain_average_op)

  return train_op
//  [python]
def smooth_l1_loss_rcnn(self, bbox_pred, bbox_targets, label, num_classes, sigma=1.0):
        '''
        :param bbox_pred: [-1, (cfgs.CLS_NUM +1) * 4]
        :param bbox_targets:[-1, (cfgs.CLS_NUM +1) * 4]
        :param label:[-1]
        :param num_classes:
        :param sigma:
        :return:
        '''

        outside_mask = tf.stop_gradient(tf.to_float(tf.greater(label, 0)))

        bbox_pred = tf.reshape(bbox_pred, [-1, num_classes, 4])
        bbox_targets = tf.reshape(bbox_targets, [-1, num_classes, 4])

        value = self._smooth_l1_loss_base(bbox_pred, bbox_targets, sigma=sigma)
        value = tf.reduce_sum(value, 2)
        value = tf.reshape(value, [-1, num_classes])

        inside_mask = tf.one_hot(tf.reshape(label, [-1, 1]),
                                 depth=num_classes, axis=1)

        inside_mask = tf.stop_gradient(
            tf.to_float(tf.reshape(inside_mask, [-1, num_classes])))

        normalizer = tf.to_float(tf.shape(bbox_pred)[0])
        bbox_loss = tf.reduce_sum(
            tf.reduce_sum(value * inside_mask, 1) * outside_mask) / normalizer

        return bbox_loss
//  [python]
def smooth_l1_loss_rcnn(self, bbox_pred, bbox_targets, label, num_classes, sigma=1.0):
        '''
        :param bbox_pred: [-1, (cfgs.CLS_NUM +1) * 4]
        :param bbox_targets:[-1, (cfgs.CLS_NUM +1) * 4]
        :param label:[-1]
        :param num_classes:
        :param sigma:
        :return:
        '''

        outside_mask = tf.stop_gradient(tf.cast(tf.greater(label, 0), tf.float32))

        bbox_pred = tf.reshape(bbox_pred, [-1, num_classes, 4])
        bbox_targets = tf.reshape(bbox_targets, [-1, num_classes, 4])

        value = self._smooth_l1_loss_base(bbox_pred, bbox_targets, sigma=sigma)
        value = tf.reduce_sum(value, 2)
        value = tf.reshape(value, [-1, num_classes])

        inside_mask = tf.one_hot(tf.reshape(label, [-1, 1]),
                                 depth=num_classes, axis=1)

        inside_mask = tf.stop_gradient(
            tf.cast(tf.reshape(inside_mask, [-1, num_classes]), tf.float32))

        normalizer = tf.cast(tf.shape(bbox_pred)[0], tf.float32)
        bbox_loss = tf.reduce_sum(
            tf.reduce_sum(value * inside_mask, 1) * outside_mask) / normalizer

        return bbox_loss
//  [python]
def load_image(image_file, image_size=None):
  """Loads an image and center-crops it to a specific size.

  Args:
    image_file: str. Image file.
    image_size: int, optional. Desired size. If provided, crops the image to
        a square and resizes it to the requested size. Defaults to None.

  Returns:
    A 4-D tensor of shape [1, image_size, image_size, 3] and dtype float32,
    with values in [0, 1].
  """
  image = tf.constant(np.uint8(load_np_image(image_file) * 255.0))
  if image_size is not None:
    # Center-crop into a square and resize to image_size
    small_side = int(min(image.shape[0], image.shape[1]))
    image = tf.image.resize_image_with_crop_or_pad(
        image, small_side, small_side)
    image = tf.image.resize_images(image, [image_size, image_size])
  image = tf.to_float(image) / 255.0

  return tf.expand_dims(image, 0)
//  [python]
def load_image(image_file, image_size=None):
  """Loads an image and center-crops it to a specific size.

  Args:
    image_file: str. Image file.
    image_size: int, optional. Desired size. If provided, crops the image to
        a square and resizes it to the requested size. Defaults to None.

  Returns:
    A 4-D tensor of shape [1, image_size, image_size, 3] and dtype float32,
    with values in [0, 1].
  """
  image = tf.constant(np.uint8(load_np_image(image_file) * 255.0))
  if image_size is not None:
    # Center-crop into a square and resize to image_size
    small_side = int(min(image.shape[0], image.shape[1]))
    image = tf.image.resize_image_with_crop_or_pad(
        image, small_side, small_side)
    image = tf.image.resize_images(image, [image_size, image_size])
  image = tf.cast(image, tf.float32) / 255.0

  return tf.expand_dims(image, 0)
//  [python]
def smooth_l1_loss_rcnn_h(self, bbox_pred, bbox_targets, label, num_classes, sigma=1.0):
        '''
        :param bbox_pred: [-1, (cfgs.CLS_NUM +1) * 4]
        :param bbox_targets:[-1, (cfgs.CLS_NUM +1) * 4]
        :param label:[-1]
        :param num_classes:
        :param sigma:
        :return:
        '''

        outside_mask = tf.stop_gradient(tf.to_float(tf.greater(label, 0)))

        bbox_pred = tf.reshape(bbox_pred, [-1, num_classes, 4])
        bbox_targets = tf.reshape(bbox_targets, [-1, num_classes, 4])

        value = self._smooth_l1_loss_base(bbox_pred,
                                          bbox_targets,
                                          sigma=sigma)
        value = tf.reduce_sum(value, 2)
        value = tf.reshape(value, [-1, num_classes])

        inside_mask = tf.one_hot(tf.reshape(label, [-1, 1]),
                                 depth=num_classes, axis=1)

        inside_mask = tf.stop_gradient(
            tf.to_float(tf.reshape(inside_mask, [-1, num_classes])))

        normalizer = tf.to_float(tf.shape(bbox_pred)[0])
        bbox_loss = tf.reduce_sum(
            tf.reduce_sum(value * inside_mask, 1) * outside_mask) / normalizer

        return bbox_loss
//  [python]
def smooth_l1_loss_rcnn_h(self, bbox_pred, bbox_targets, label, num_classes, sigma=1.0):
        '''
        :param bbox_pred: [-1, (cfgs.CLS_NUM +1) * 4]
        :param bbox_targets:[-1, (cfgs.CLS_NUM +1) * 4]
        :param label:[-1]
        :param num_classes:
        :param sigma:
        :return:
        '''

        outside_mask = tf.stop_gradient(tf.cast(tf.greater(label, 0), tf.float32))

        bbox_pred = tf.reshape(bbox_pred, [-1, num_classes, 4])
        bbox_targets = tf.reshape(bbox_targets, [-1, num_classes, 4])

        value = self._smooth_l1_loss_base(bbox_pred,
                                          bbox_targets,
                                          sigma=sigma)
        value = tf.reduce_sum(value, 2)
        value = tf.reshape(value, [-1, num_classes])

        inside_mask = tf.one_hot(tf.reshape(label, [-1, 1]),
                                 depth=num_classes, axis=1)

        inside_mask = tf.stop_gradient(
            tf.cast(tf.reshape(inside_mask, [-1, num_classes]), tf.float32))

        normalizer = tf.cast(tf.shape(bbox_pred)[0], tf.float32)
        bbox_loss = tf.reduce_sum(
            tf.reduce_sum(value * inside_mask, 1) * outside_mask) / normalizer

        return bbox_loss
//  [python]
def center_crop_resize_image(image, image_size):
  """Center-crop into a square and resize to image_size.

  Args:
    image: A 3-D image `Tensor`.
    image_size: int, Desired size. Crops the image to a square and resizes it
      to the requested size.

  Returns:
    A 4-D tensor of shape [1, image_size, image_size, 3] and dtype float32,
    with values in [0, 1].
  """
  shape = tf.shape(image)
  small_side = tf.minimum(shape[0], shape[1])
  image = tf.image.resize_image_with_crop_or_pad(image, small_side, small_side)
  image = tf.to_float(image) / 255.0

  image = tf.image.resize_images(image, tf.constant([image_size, image_size]))

  return tf.expand_dims(image, 0)
//  [python]
def center_crop_resize_image(image, image_size):
  """Center-crop into a square and resize to image_size.

  Args:
    image: A 3-D image `Tensor`.
    image_size: int, Desired size. Crops the image to a square and resizes it
      to the requested size.

  Returns:
    A 4-D tensor of shape [1, image_size, image_size, 3] and dtype float32,
    with values in [0, 1].
  """
  shape = tf.shape(image)
  small_side = tf.minimum(shape[0], shape[1])
  image = tf.image.resize_image_with_crop_or_pad(image, small_side, small_side)
  image = tf.cast(image, tf.float32) / 255.0

  image = tf.image.resize_images(image, tf.constant([image_size, image_size]))

  return tf.expand_dims(image, 0)
//  [python]
def smooth_l1_loss_rcnn_r(self, bbox_pred, bbox_targets, label, num_classes, sigma=1.0):
        '''
        :param bbox_pred: [-1, (cfgs.CLS_NUM +1) * 5]
        :param bbox_targets:[-1, (cfgs.CLS_NUM +1) * 5]
        :param label:[-1]
        :param num_classes:
        :param sigma:
        :return:
        '''

        outside_mask = tf.stop_gradient(tf.to_float(tf.greater(label, 0)))

        bbox_pred = tf.reshape(bbox_pred, [-1, num_classes, 5])
        bbox_targets = tf.reshape(bbox_targets, [-1, num_classes, 5])

        value = self._smooth_l1_loss_base(bbox_pred, bbox_targets, sigma=sigma)
        value = tf.reduce_sum(value, 2)
        value = tf.reshape(value, [-1, num_classes])

        inside_mask = tf.one_hot(tf.reshape(label, [-1, 1]),
                                 depth=num_classes, axis=1)

        inside_mask = tf.stop_gradient(
            tf.to_float(tf.reshape(inside_mask, [-1, num_classes])))

        normalizer = tf.to_float(tf.shape(bbox_pred)[0])
        bbox_loss = tf.reduce_sum(
            tf.reduce_sum(value * inside_mask, 1) * outside_mask) / normalizer

        return bbox_loss
//  [python]
def smooth_l1_loss_rcnn_r(self, bbox_pred, bbox_targets, label, num_classes, sigma=1.0):
        '''
        :param bbox_pred: [-1, (cfgs.CLS_NUM +1) * 5]
        :param bbox_targets:[-1, (cfgs.CLS_NUM +1) * 5]
        :param label:[-1]
        :param num_classes:
        :param sigma:
        :return:
        '''

        outside_mask = tf.stop_gradient(tf.cast(tf.greater(label, 0), tf.float32))

        bbox_pred = tf.reshape(bbox_pred, [-1, num_classes, 5])
        bbox_targets = tf.reshape(bbox_targets, [-1, num_classes, 5])

        value = self._smooth_l1_loss_base(bbox_pred, bbox_targets, sigma=sigma)
        value = tf.reduce_sum(value, 2)
        value = tf.reshape(value, [-1, num_classes])

        inside_mask = tf.one_hot(tf.reshape(label, [-1, 1]),
                                 depth=num_classes, axis=1)

        inside_mask = tf.stop_gradient(
            tf.cast(tf.reshape(inside_mask, [-1, num_classes]), tf.float32))

        normalizer = tf.cast(tf.shape(bbox_pred)[0], tf.float32)
        bbox_loss = tf.reduce_sum(
            tf.reduce_sum(value * inside_mask, 1) * outside_mask) / normalizer

        return bbox_loss
//  [python]
def resize_image(image, image_size):
  """Resize input image preserving the original aspect ratio.

  Args:
    image: A 3-D image `Tensor`.
    image_size: int, desired size of the smallest size of image after resize.

  Returns:
    A 4-D tensor of shape [1, image_size, image_size, 3] and dtype float32,
    with values in [0, 1].
  """
  image = _aspect_preserving_resize(image, image_size)
  image = tf.to_float(image) / 255.0

  return tf.expand_dims(image, 0)
//  [python]
def resize_image(image, image_size):
  """Resize input image preserving the original aspect ratio.

  Args:
    image: A 3-D image `Tensor`.
    image_size: int, desired size of the smallest size of image after resize.

  Returns:
    A 4-D tensor of shape [1, image_size, image_size, 3] and dtype float32,
    with values in [0, 1].
  """
  image = _aspect_preserving_resize(image, image_size)
  image = tf.cast(image, tf.float32) / 255.0

  return tf.expand_dims(image, 0)
//  [python]
def resize_video_frames(images, size):
  return [tf.to_int64(tf.image.resize_images(
      image, [size, size], tf.image.ResizeMethod.BILINEAR)) for image in images]
//  [python]
def resize_video_frames(images, size):
  return [tf.cast(tf.image.resize(
      image, [size, size], tf.image.ResizeMethod.BILINEAR), tf.int64) for image in images]
//  [python]
def exclude_ignore_index(self, logits, labels):
        """Ignore certain index."""
        logits = tf.transpose(logits, [0, 2, 3, 1])
        if zeus.is_gpu_device():
            indices = tf.where(tf.not_equal(labels, self.ignore_index))
            labels = tf.cast(tf.gather_nd(labels, indices), tf.int32)
            logits = tf.gather_nd(logits, indices)
            return logits, labels, 1.0
        elif zeus.is_npu_device():
            weights = tf.not_equal(labels, self.ignore_index)
            labels = tf.multiply(labels, tf.cast(weights, labels.dtype))
            return logits, labels, tf.to_float(weights)
//  [python]
def exclude_ignore_index(self, logits, labels):
        """Ignore certain index."""
        logits = tf.transpose(logits, [0, 2, 3, 1])
        if zeus.is_gpu_device():
            indices = tf.where(tf.not_equal(labels, self.ignore_index))
            labels = tf.cast(tf.gather_nd(labels, indices), tf.int32)
            logits = tf.gather_nd(logits, indices)
            return logits, labels, 1.0
        elif zeus.is_npu_device():
            weights = tf.not_equal(labels, self.ignore_index)
            labels = tf.multiply(labels, tf.cast(weights, labels.dtype))
            return logits, labels, tf.cast(weights, tf.float32)
//  [python]
def bilstm_reader(placeholders, vocab_size, emb_dim, drop_keep_prob=1.0):
    # [batch_size, max_seq_length]
    sentences = placeholders['sentences']

    # [batch_size, candidate_size]
    targets = tf.to_float(placeholders['sentence_labels'])

    with tf.variable_scope("embeddings"):
        embeddings = tf.get_variable("word_embeddings", [vocab_size, emb_dim], dtype=tf.float32)

    with tf.variable_scope("embedders") as varscope:
        sentences_embedded = tf.nn.embedding_lookup(embeddings, sentences)

    with tf.variable_scope("bilstm_reader") as varscope1:

        # states: (c_fw, h_fw), (c_bw, h_bw)
        outputs, states = reader(sentences_embedded, placeholders['sentences_lengths'], emb_dim,
                                scope=varscope1, drop_keep_prob=drop_keep_prob)

        # concat fw and bw outputs
        output = tf.concat(1, [states[0][1], states[1][1]])

    scores = tf.contrib.layers.linear(output, 2)  # we don't strictly need this as we've only got 2 targets
    # add non-linearity
    scores = tf.nn.tanh(scores)
    loss = tf.nn.softmax_cross_entropy_with_logits(scores, targets)
    predict = tf.nn.softmax(scores)

    predictions = tf.argmax(predict, axis=1)
    true_vals = tf.argmax(targets, axis=1)
    accuracy = tf.reduce_mean(tf.cast(tf.equal(predictions, true_vals), tf.float32))

    saver = tf.train.Saver()

    return scores, loss, predict, accuracy, saver
//  [python]
def bilstm_reader(placeholders, vocab_size, emb_dim, drop_keep_prob=1.0):
    # [batch_size, max_seq_length]
    sentences = placeholders['sentences']

    # [batch_size, candidate_size]
    targets = tf.cast(placeholders['sentence_labels'], tf.float32)

    with tf.variable_scope("embeddings"):
        embeddings = tf.get_variable("word_embeddings", [vocab_size, emb_dim], dtype=tf.float32)

    with tf.variable_scope("embedders") as varscope:
        sentences_embedded = tf.nn.embedding_lookup(embeddings, sentences)

    with tf.variable_scope("bilstm_reader") as varscope1:

        # states: (c_fw, h_fw), (c_bw, h_bw)
        outputs, states = reader(sentences_embedded, placeholders['sentences_lengths'], emb_dim,
                                scope=varscope1, drop_keep_prob=drop_keep_prob)

        # concat fw and bw outputs
        output = tf.concat(1, [states[0][1], states[1][1]])

    scores = tf.contrib.layers.linear(output, 2)  # we don't strictly need this as we've only got 2 targets
    # add non-linearity
    scores = tf.nn.tanh(scores)
    loss = tf.nn.softmax_cross_entropy_with_logits(scores, targets)
    predict = tf.nn.softmax(scores)

    predictions = tf.argmax(predict, axis=1)
    true_vals = tf.argmax(targets, axis=1)
    accuracy = tf.reduce_mean(tf.cast(tf.equal(predictions, true_vals), tf.float32))

    saver = tf.train.Saver()

    return scores, loss, predict, accuracy, saver
//  [python]
def colorize_img(value, vmin=None, vmax=None, cmap=None):
    """
    A utility function for TensorFlow that maps a grayscale image to a matplotlib colormap for use with TensorBoard image summaries.
    By default it will normalize the input value to the range 0..1 before mapping to a grayscale colormap.
    Arguments:
      - value: 4D Tensor of shape [batch_size,height, width,1]
      - vmin: the minimum value of the range used for normalization. (Default: value minimum)
      - vmax: the maximum value of the range used for normalization. (Default: value maximum)
      - cmap: a valid cmap named for use with matplotlib's 'get_cmap'.(Default: 'gray')
    
    Returns a 3D tensor of shape [batch_size,height, width,3].
    """

    # normalize
    vmin = tf.reduce_min(value) if vmin is None else vmin
    vmax = tf.reduce_max(value) if vmax is None else vmax
    value = (value - vmin) / (vmax - vmin) # vmin..vmax

    # quantize
    indices = tf.to_int32(tf.round(value[:,:,:,0]*255))

    # gather
    color_map = cm.get_cmap(cmap if cmap is not None else 'gray')
    colors = color_map(np.arange(256))[:,:3]
    colors = tf.constant(colors, dtype=tf.float32)
    value = tf.gather(colors, indices)
    return value
//  [python]
def colorize_img(value, vmin=None, vmax=None, cmap=None):
    """
    A utility function for TensorFlow that maps a grayscale image to a matplotlib colormap for use with TensorBoard image summaries.
    By default it will normalize the input value to the range 0..1 before mapping to a grayscale colormap.
    Arguments:
      - value: 4D Tensor of shape [batch_size,height, width,1]
      - vmin: the minimum value of the range used for normalization. (Default: value minimum)
      - vmax: the maximum value of the range used for normalization. (Default: value maximum)
      - cmap: a valid cmap named for use with matplotlib's 'get_cmap'.(Default: 'gray')
    
    Returns a 3D tensor of shape [batch_size,height, width,3].
    """

    # normalize
    vmin = tf.reduce_min(value) if vmin is None else vmin
    vmax = tf.reduce_max(value) if vmax is None else vmax
    value = (value - vmin) / (vmax - vmin) # vmin..vmax

    # quantize
    indices = tf.cast(tf.round(value[:,:,:,0]*255), tf.int32)

    # gather
    color_map = cm.get_cmap(cmap if cmap is not None else 'gray')
    colors = color_map(np.arange(256))[:,:3]
    colors = tf.constant(colors, dtype=tf.float32)
    value = tf.gather(colors, indices)
    return value
//  [python]
def get_input_fn(self,
                   partition,
                   mode,
                   batch_size):
    """See `data.Provider` get_input_fn."""

    del partition
    def input_fn(params=None):
      """Input_fn to return."""

      del params  # Unused.

      np.random.seed(self._seed)
      if mode == tf.estimator.ModeKeys.EVAL:
        np.random.seed(self._seed + 1)

      images = tf.to_float(
          tf.convert_to_tensor(
              np.random.rand(self._num_examples, *self._shape())))
      labels = tf.convert_to_tensor(
          np.random.randint(0, high=2, size=(self._num_examples, 1)))
      dataset = tf.data.Dataset.from_tensor_slices(({"x": images}, labels))
      if mode == tf.estimator.ModeKeys.TRAIN:
        dataset = dataset.repeat()
      dataset = dataset.batch(batch_size)
      iterator = dataset.make_one_shot_iterator()
      return iterator.get_next()

    return input_fn
//  [python]
def get_input_fn(self,
                   partition,
                   mode,
                   batch_size):
    """See `data.Provider` get_input_fn."""

    del partition
    def input_fn(params=None):
      """Input_fn to return."""

      del params  # Unused.

      np.random.seed(self._seed)
      if mode == tf.estimator.ModeKeys.EVAL:
        np.random.seed(self._seed + 1)

      images = tf.cast(
          tf.convert_to_tensor(
              np.random.rand(self._num_examples, *self._shape())), tf.float32)
      labels = tf.convert_to_tensor(
          np.random.randint(0, high=2, size=(self._num_examples, 1)))
      dataset = tf.data.Dataset.from_tensor_slices(({"x": images}, labels))
      if mode == tf.estimator.ModeKeys.TRAIN:
        dataset = dataset.repeat()
      dataset = dataset.batch(batch_size)
      iterator = dataset.make_one_shot_iterator()
      return iterator.get_next()

    return input_fn