jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def cross_entropy_seq_with_mask(logits, target_seqs, input_mask, return_details=False, name=None):
    """Returns the expression of cross-entropy of two sequences, implement
    softmax internally. Normally be used for Dynamic RNN outputs.

    Parameters
    -----------
    logits : network identity outputs
        2D tensor, ``network.outputs``, [batch_size, number of output units].
    target_seqs : int of tensor, like word ID.
        [batch_size, ?]
    input_mask : the mask to compute loss
        The same size with target_seqs, normally 0 and 1.
    return_details : boolean
        - If False (default), only returns the loss.
        - If True, returns the loss, losses, weights and targets (reshape to one vetcor).

    Examples
    --------
    - see Image Captioning Example.
    """
    targets = tf.reshape(target_seqs, [-1])   # to one vector
    weights = tf.cast(tf.reshape(input_mask, [-1]), tf.float32)   # to one vector like targets
    losses = tf.nn.sparse_softmax_cross_entropy_with_logits(logits=logits, labels=targets, name=name) * weights
    #losses = tf.reduce_mean(tf.nn.sparse_softmax_cross_entropy_with_logits(logits=logits, labels=targets, name=name)) # for TF1.0 and others

    try: ## TF1.0
        loss = tf.divide(tf.reduce_sum(losses),   # loss from mask. reduce_sum before element-wise mul with mask !!
                        tf.reduce_sum(weights),
                        name="seq_loss_with_mask")
    except: ## TF0.12
        loss = tf.div(tf.reduce_sum(losses),   # loss from mask. reduce_sum before element-wise mul with mask !!
                        tf.reduce_sum(weights),
                        name="seq_loss_with_mask")
    if return_details:
        return loss, losses, weights, targets
    else:
        return loss
//  [python]
def Adv_graph(sess, pha):
  u = Genc.build(train_images, pha['is_training_a'])
  Y_score = Adv.build(u, N_CLASS, pha['is_training_a'])
  Y_hat = tf.nn.softmax(Y_score)
  Y_pred = tf.argmax(Y_hat, 1)
  Y_label = tf.to_float(tf.one_hot(train_labels, N_CLASS))

  cross_entropy = tf.nn.softmax_cross_entropy_with_logits(logits = Y_score, labels = Y_label) #shape [None,2]
  loss_op = tf.reduce_mean(cross_entropy)
  correct_prediction = tf.equal(tf.argmax(Y_hat, 1), tf.argmax(Y_label, 1))
  acc_op = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))

  update_op = tf.train.MomentumOptimizer(LR, MOMENTUM).minimize(loss_op, var_list=Adv.vars)

  return loss_op, acc_op, update_op, Y_pred
//  [python]
def Adv_graph(sess, pha):
  u = Genc.build(train_images, pha['is_training_a'])
  Y_score = Adv.build(u, N_CLASS, pha['is_training_a'])
  Y_hat = tf.nn.softmax(Y_score)
  Y_pred = tf.argmax(Y_hat, 1)
  Y_label = tf.cast(tf.one_hot(train_labels, N_CLASS), tf.float32)

  cross_entropy = tf.nn.softmax_cross_entropy_with_logits(logits = Y_score, labels = Y_label) #shape [None,2]
  loss_op = tf.reduce_mean(cross_entropy)
  correct_prediction = tf.equal(tf.argmax(Y_hat, 1), tf.argmax(Y_label, 1))
  acc_op = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))

  update_op = tf.train.MomentumOptimizer(LR, MOMENTUM).minimize(loss_op, var_list=Adv.vars)

  return loss_op, acc_op, update_op, Y_pred
//  [python]
def get_candidate_labels(self, candidate_starts, candidate_ends, labeled_starts, labeled_ends, labels):
        same_start = tf.equal(tf.expand_dims(labeled_starts, 1),
                              tf.expand_dims(candidate_starts, 0))  # [num_labeled, num_candidates]
        same_end = tf.equal(tf.expand_dims(labeled_ends, 1),
                            tf.expand_dims(candidate_ends, 0))  # [num_labeled, num_candidates]
        same_span = tf.to_int32(tf.logical_and(same_start, same_end))  # [num_labeled, num_candidates]
        is_gold_span = tf.reduce_sum(same_span, axis=0)
        candidate_labels = tf.matmul(tf.expand_dims(labels, 0), same_span)  # [1, num_candidates]
        candidate_labels = tf.squeeze(candidate_labels, 0)  # [num_candidates]
        return candidate_labels, is_gold_span
//  [python]
def get_candidate_labels(self, candidate_starts, candidate_ends, labeled_starts, labeled_ends, labels):
        same_start = tf.equal(tf.expand_dims(labeled_starts, 1),
                              tf.expand_dims(candidate_starts, 0))  # [num_labeled, num_candidates]
        same_end = tf.equal(tf.expand_dims(labeled_ends, 1),
                            tf.expand_dims(candidate_ends, 0))  # [num_labeled, num_candidates]
        same_span = tf.cast(tf.logical_and(same_start, same_end), tf.int32)  # [num_labeled, num_candidates]
        is_gold_span = tf.reduce_sum(same_span, axis=0)
        candidate_labels = tf.matmul(tf.expand_dims(labels, 0), same_span)  # [1, num_candidates]
        candidate_labels = tf.squeeze(candidate_labels, 0)  # [num_candidates]
        return candidate_labels, is_gold_span
//  [python]
def add_timing_signal(self, x, min_timescale=1.0, max_timescale=1.0e4):
        """Adds a bunch of sinusoids of different frequencies to a Tensor.
        Each channel of the input Tensor is incremented by a sinusoid of a
        different frequency and phase.
        This allows attention to learn to use absolute and relative positions.
        Timing signals should be added to some precursors of both the query and the
        memory inputs to attention.
        The use of relative position is possible because sin(x+y) and cos(x+y) can
        be experessed in terms of y, sin(x) and cos(x).
        In particular, we use a geometric sequence of timescales starting with
        min_timescale and ending with max_timescale.  The number of different
        timescales is equal to channels / 2. For each timescale, we
        generate the two sinusoidal signals sin(timestep/timescale) and
        cos(timestep/timescale).  All of these sinusoids are concatenated in
        the channels dimension.
        Args:
            x: a Tensor with shape [batch, length, channels]
            min_timescale: a float
            max_timescale: a float
        Returns:
            a Tensor the same shape as x.
        """
        with tf.name_scope("add_timing_signal", values=[x]):
            length = tf.shape(x)[1]
            channels = tf.shape(x)[2]
            position = tf.to_float(tf.range(length))
            num_timescales = channels // 2

            log_timescale_increment = (
                math.log(float(max_timescale) / float(min_timescale)) /
                (tf.to_float(num_timescales) - 1)
            )
            inv_timescales = min_timescale * tf.exp(
                tf.to_float(tf.range(num_timescales)) * -log_timescale_increment
            )

            scaled_time = (tf.expand_dims(position, 1) *
                           tf.expand_dims(inv_timescales, 0))
            signal = tf.concat([tf.sin(scaled_time), tf.cos(scaled_time)], axis=1)
            signal = tf.pad(signal, [[0, 0], [0, tf.mod(channels, 2)]])
            signal = tf.reshape(signal, [1, length, channels])

            return x + signal
//  [python]
def add_timing_signal(self, x, min_timescale=1.0, max_timescale=1.0e4):
        """Adds a bunch of sinusoids of different frequencies to a Tensor.
        Each channel of the input Tensor is incremented by a sinusoid of a
        different frequency and phase.
        This allows attention to learn to use absolute and relative positions.
        Timing signals should be added to some precursors of both the query and the
        memory inputs to attention.
        The use of relative position is possible because sin(x+y) and cos(x+y) can
        be experessed in terms of y, sin(x) and cos(x).
        In particular, we use a geometric sequence of timescales starting with
        min_timescale and ending with max_timescale.  The number of different
        timescales is equal to channels / 2. For each timescale, we
        generate the two sinusoidal signals sin(timestep/timescale) and
        cos(timestep/timescale).  All of these sinusoids are concatenated in
        the channels dimension.
        Args:
            x: a Tensor with shape [batch, length, channels]
            min_timescale: a float
            max_timescale: a float
        Returns:
            a Tensor the same shape as x.
        """
        with tf.name_scope("add_timing_signal", values=[x]):
            length = tf.shape(x)[1]
            channels = tf.shape(x)[2]
            position = tf.cast(tf.range(length), tf.float32)
            num_timescales = channels // 2

            log_timescale_increment = (
                math.log(float(max_timescale) / float(min_timescale)) /
                (tf.cast(num_timescales, tf.float32) - 1)
            )
            inv_timescales = min_timescale * tf.exp(
                tf.cast(tf.range(num_timescales), tf.float32) * -log_timescale_increment
            )

            scaled_time = (tf.expand_dims(position, 1) *
                           tf.expand_dims(inv_timescales, 0))
            signal = tf.concat([tf.sin(scaled_time), tf.cos(scaled_time)], axis=1)
            signal = tf.pad(signal, [[0, 0], [0, tf.mod(channels, 2)]])
            signal = tf.reshape(signal, [1, length, channels])

            return x + signal
//  [python]
def G_adv_graph(sess, phg):
  real_labels = train_labels * 2 -1
  fake_labels = -real_labels

  u = Genc.build(train_images, phg['is_training_g'])

  Y_score = Adv.build(u, N_CLASS, phg['is_training_g'])
  Y_hat = tf.nn.softmax(Y_score)
  Y_pred = tf.argmax(Y_hat, 1)
  Y_label = tf.to_float(tf.one_hot(train_labels, N_CLASS))
  cross_entropy = tf.nn.softmax_cross_entropy_with_logits(logits = Y_score, labels = Y_label) #shape [None,2]
  adv_loss = tf.reduce_mean(cross_entropy)
  correct_prediction = tf.equal(tf.argmax(Y_hat, 1), tf.argmax(Y_label, 1))
  acc_op = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))


  real_images = Gdec.build(u, real_labels, phg['is_training_g'])
  fake_images = Gdec.build(u, fake_labels, phg['is_training_g'])
  fake_gan_logit, fake_cls_logit = D.build(fake_images, phg['is_training_g'])

  gan_loss = -tf.reduce_mean(fake_gan_logit)
  cls_loss = tf.losses.sigmoid_cross_entropy(tf.expand_dims((1-train_labels),axis=-1), fake_cls_logit)
  rec_loss = tf.losses.absolute_difference(train_images, real_images)
  reg_loss = tf.reduce_sum(Genc.reg_loss+Gdec.reg_loss)

  final_loss = gan_loss + 10.0*cls_loss + 100.0*rec_loss + reg_loss - adv_loss
  update_op = tf.train.AdamOptimizer(phg['lr_g'], beta1=0.5, name='second_adam').minimize(final_loss, var_list=Genc.vars+Gdec.vars)

  return u, final_loss, acc_op, update_op
//  [python]
def G_adv_graph(sess, phg):
  real_labels = train_labels * 2 -1
  fake_labels = -real_labels

  u = Genc.build(train_images, phg['is_training_g'])

  Y_score = Adv.build(u, N_CLASS, phg['is_training_g'])
  Y_hat = tf.nn.softmax(Y_score)
  Y_pred = tf.argmax(Y_hat, 1)
  Y_label = tf.cast(tf.one_hot(train_labels, N_CLASS), tf.float32)
  cross_entropy = tf.nn.softmax_cross_entropy_with_logits(logits = Y_score, labels = Y_label) #shape [None,2]
  adv_loss = tf.reduce_mean(cross_entropy)
  correct_prediction = tf.equal(tf.argmax(Y_hat, 1), tf.argmax(Y_label, 1))
  acc_op = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))


  real_images = Gdec.build(u, real_labels, phg['is_training_g'])
  fake_images = Gdec.build(u, fake_labels, phg['is_training_g'])
  fake_gan_logit, fake_cls_logit = D.build(fake_images, phg['is_training_g'])

  gan_loss = -tf.reduce_mean(fake_gan_logit)
  cls_loss = tf.losses.sigmoid_cross_entropy(tf.expand_dims((1-train_labels),axis=-1), fake_cls_logit)
  rec_loss = tf.losses.absolute_difference(train_images, real_images)
  reg_loss = tf.reduce_sum(Genc.reg_loss+Gdec.reg_loss)

  final_loss = gan_loss + 10.0*cls_loss + 100.0*rec_loss + reg_loss - adv_loss
  update_op = tf.train.AdamOptimizer(phg['lr_g'], beta1=0.5, name='second_adam').minimize(final_loss, var_list=Genc.vars+Gdec.vars)

  return u, final_loss, acc_op, update_op
//  [python]
def bucket_distance(self, distances):
        """
        Places the given values (designed for distances) into 10 semi-logscale buckets:
        [0, 1, 2, 3, 4, 5-7, 8-15, 16-31, 32-63, 64+].
        """
        logspace_idx = tf.to_int32(tf.floor(tf.log(tf.to_float(distances)) / math.log(2))) + 3
        use_identity = tf.to_int32(distances <= 4)
        combined_idx = use_identity * distances + (1 - use_identity) * logspace_idx
        return tf.clip_by_value(combined_idx, 0, 9)
//  [python]
def bucket_distance(self, distances):
        """
        Places the given values (designed for distances) into 10 semi-logscale buckets:
        [0, 1, 2, 3, 4, 5-7, 8-15, 16-31, 32-63, 64+].
        """
        logspace_idx = tf.cast(tf.floor(tf.log(tf.cast(distances, tf.float32)) / math.log(2)), tf.int32) + 3
        use_identity = tf.cast(distances <= 4, tf.int32)
        combined_idx = use_identity * distances + (1 - use_identity) * logspace_idx
        return tf.clip_by_value(combined_idx, 0, 9)
//  [python]
def pad_to_multiple(tensor, multiple):
  """Returns the tensor zero padded to the specified multiple.

  Appends 0s to the end of the first and second dimension (height and width) of
  the tensor until both dimensions are a multiple of the input argument
  'multiple'. E.g. given an input tensor of shape [1, 3, 5, 1] and an input
  multiple of 4, PadToMultiple will append 0s so that the resulting tensor will
  be of shape [1, 4, 8, 1].

  Args:
    tensor: rank 4 float32 tensor, where
            tensor -> [batch_size, height, width, channels].
    multiple: the multiple to pad to.

  Returns:
    padded_tensor: the tensor zero padded to the specified multiple.
  """
  if multiple == 1:
    return tensor

  tensor_shape = tensor.get_shape()
  batch_size = static_shape.get_batch_size(tensor_shape)
  tensor_height = static_shape.get_height(tensor_shape)
  tensor_width = static_shape.get_width(tensor_shape)
  tensor_depth = static_shape.get_depth(tensor_shape)

  if batch_size is None:
    batch_size = tf.shape(tensor)[0]

  if tensor_height is None:
    tensor_height = tf.shape(tensor)[1]
    padded_tensor_height = tf.to_int32(
        tf.ceil(tf.to_float(tensor_height) / tf.to_float(multiple))) * multiple
  else:
    padded_tensor_height = int(
        math.ceil(float(tensor_height) / multiple) * multiple)

  if tensor_width is None:
    tensor_width = tf.shape(tensor)[2]
    padded_tensor_width = tf.to_int32(
        tf.ceil(tf.to_float(tensor_width) / tf.to_float(multiple))) * multiple
  else:
    padded_tensor_width = int(
        math.ceil(float(tensor_width) / multiple) * multiple)

  if tensor_depth is None:
    tensor_depth = tf.shape(tensor)[3]

  # Use tf.concat instead of tf.pad to preserve static shape
  if padded_tensor_height != tensor_height:
    height_pad = tf.zeros([
        batch_size, padded_tensor_height - tensor_height, tensor_width,
        tensor_depth
    ])
    tensor = tf.concat([tensor, height_pad], 1)
  if padded_tensor_width != tensor_width:
    width_pad = tf.zeros([
        batch_size, padded_tensor_height, padded_tensor_width - tensor_width,
        tensor_depth
    ])
    tensor = tf.concat([tensor, width_pad], 2)

  return tensor
//  [python]
def pad_to_multiple(tensor, multiple):
  """Returns the tensor zero padded to the specified multiple.

  Appends 0s to the end of the first and second dimension (height and width) of
  the tensor until both dimensions are a multiple of the input argument
  'multiple'. E.g. given an input tensor of shape [1, 3, 5, 1] and an input
  multiple of 4, PadToMultiple will append 0s so that the resulting tensor will
  be of shape [1, 4, 8, 1].

  Args:
    tensor: rank 4 float32 tensor, where
            tensor -> [batch_size, height, width, channels].
    multiple: the multiple to pad to.

  Returns:
    padded_tensor: the tensor zero padded to the specified multiple.
  """
  if multiple == 1:
    return tensor

  tensor_shape = tensor.get_shape()
  batch_size = static_shape.get_batch_size(tensor_shape)
  tensor_height = static_shape.get_height(tensor_shape)
  tensor_width = static_shape.get_width(tensor_shape)
  tensor_depth = static_shape.get_depth(tensor_shape)

  if batch_size is None:
    batch_size = tf.shape(tensor)[0]

  if tensor_height is None:
    tensor_height = tf.shape(tensor)[1]
    padded_tensor_height = tf.cast(
        tf.ceil(tf.to_float(tensor_height) / tf.to_float(multiple))) * multiple
  else:
    padded_tensor_height = int(
        math.ceil(float(tensor_height) / multiple) * multiple)

  if tensor_width is None:
    tensor_width = tf.shape(tensor)[2]
    padded_tensor_width = tf.cast(
        tf.ceil(tf.to_float(tensor_width) / tf.to_float(multiple))) * multiple
  else:
    padded_tensor_width = int(
        math.ceil(float(tensor_width) / multiple) * multiple)

  if tensor_depth is None:
    tensor_depth = tf.shape(tensor)[3]

  # Use tf.concat instead of tf.pad to preserve static shape
  if padded_tensor_height != tensor_height:
    height_pad = tf.zeros([
        batch_size, padded_tensor_height - tensor_height, tensor_width,
        tensor_depth
    ])
    tensor = tf.concat([tensor, height_pad], 1)
  if padded_tensor_width != tensor_width:
    width_pad = tf.zeros([
        batch_size, padded_tensor_height, padded_tensor_width - tensor_width,
        tensor_depth
    ])
    tensor = tf.concat([tensor, width_pad], 2)

  return tensor
//  [python]
def make_outer_masks(self, outer_masks, input_pianorolls):
    """Returns outer masks, if all zeros created by completion masking."""
    outer_masks = tf.to_float(outer_masks)
    # If outer_masks come in as all zeros, it means there's no masking,
    # which also means nothing will be generated. In this case, use
    # completion mask to make new outer masks.
    outer_masks = tf.cond(
        tf.reduce_all(tf.equal(outer_masks, 0)),
        lambda: make_completion_masks(input_pianorolls),
        lambda: outer_masks)
    return outer_masks
//  [python]
def make_outer_masks(self, outer_masks, input_pianorolls):
    """Returns outer masks, if all zeros created by completion masking."""
    outer_masks = tf.cast(outer_masks, tf.float32)
    # If outer_masks come in as all zeros, it means there's no masking,
    # which also means nothing will be generated. In this case, use
    # completion mask to make new outer masks.
    outer_masks = tf.cond(
        tf.reduce_all(tf.equal(outer_masks, 0)),
        lambda: make_completion_masks(input_pianorolls),
        lambda: outer_masks)
    return outer_masks
//  [python]
def coarse_to_fine_pruning(self, top_span_emb, top_span_mention_scores, c):
        k = util.shape(top_span_emb, 0)
        top_span_range = tf.range(k)  # [k]
        antecedent_offsets = tf.expand_dims(top_span_range, 1) - tf.expand_dims(top_span_range, 0)  # [k, k]
        antecedents_mask = antecedent_offsets >= 1  # [k, k]
        fast_antecedent_scores = tf.expand_dims(top_span_mention_scores, 1) + tf.expand_dims(top_span_mention_scores,
                                                                                             0)  # [k, k]
        fast_antecedent_scores += tf.log(tf.to_float(antecedents_mask))  # [k, k]
        fast_antecedent_scores += self.get_fast_antecedent_scores(top_span_emb)  # [k, k]

        _, top_antecedents = tf.nn.top_k(fast_antecedent_scores, c, sorted=False)  # [k, c]
        top_antecedents_mask = util.batch_gather(antecedents_mask, top_antecedents)  # [k, c]
        top_fast_antecedent_scores = util.batch_gather(fast_antecedent_scores, top_antecedents)  # [k, c]
        top_antecedent_offsets = util.batch_gather(antecedent_offsets, top_antecedents)  # [k, c]
        return top_antecedents, top_antecedents_mask, top_fast_antecedent_scores, top_antecedent_offsets
//  [python]
def coarse_to_fine_pruning(self, top_span_emb, top_span_mention_scores, c):
        k = util.shape(top_span_emb, 0)
        top_span_range = tf.range(k)  # [k]
        antecedent_offsets = tf.expand_dims(top_span_range, 1) - tf.expand_dims(top_span_range, 0)  # [k, k]
        antecedents_mask = antecedent_offsets >= 1  # [k, k]
        fast_antecedent_scores = tf.expand_dims(top_span_mention_scores, 1) + tf.expand_dims(top_span_mention_scores,
                                                                                             0)  # [k, k]
        fast_antecedent_scores += tf.log(tf.cast(antecedents_mask, dtype=tf.float32))  # [k, k]
        fast_antecedent_scores += self.get_fast_antecedent_scores(top_span_emb)  # [k, k]

        _, top_antecedents = tf.nn.top_k(fast_antecedent_scores, c, sorted=False)  # [k, c]
        top_antecedents_mask = util.batch_gather(antecedents_mask, top_antecedents)  # [k, c]
        top_fast_antecedent_scores = util.batch_gather(fast_antecedent_scores, top_antecedents)  # [k, c]
        top_antecedent_offsets = util.batch_gather(antecedent_offsets, top_antecedents)  # [k, c]
        return top_antecedents, top_antecedents_mask, top_fast_antecedent_scores, top_antecedent_offsets
//  [python]
def predict(self, pianorolls, masks):
    """Evalutes the model once and returns predictions."""
    direct_inputs = dict(
        pianorolls=pianorolls, masks=masks,
        lengths=tf.to_float([tf.shape(pianorolls)[1]]))

    model = lib_graph.build_graph(
        is_training=False,
        hparams=self.hparams,
        direct_inputs=direct_inputs,
        use_placeholders=False)
    self.logits = model.logits
    return self.logits
//  [python]
def predict(self, pianorolls, masks):
    """Evalutes the model once and returns predictions."""
    direct_inputs = dict(
        pianorolls=pianorolls, masks=masks,
        lengths=tf.cast([tf.shape(pianorolls)[1]], tf.float32))

    model = lib_graph.build_graph(
        is_training=False,
        hparams=self.hparams,
        direct_inputs=direct_inputs,
        use_placeholders=False)
    self.logits = model.logits
    return self.logits
//  [python]
def make_completion_masks(pianorolls, outer_masks=1.):
  pianorolls = tf.to_float(pianorolls)
  masks = tf.reduce_all(tf.equal(pianorolls, 0), axis=2, keep_dims=True)
  inner_masks = tf.to_float(masks) + 0 * pianorolls
  return inner_masks * outer_masks
//  [python]
def make_completion_masks(pianorolls, outer_masks=1.):
  pianorolls = tf.cast(pianorolls, tf.float32)
  masks = tf.reduce_all(tf.equal(pianorolls, 0), axis=2, keep_dims=True)
  inner_masks = tf.cast(masks, tf.float32) + 0 * pianorolls
  return inner_masks * outer_masks
//  [python]
def make_bernoulli_masks(shape, pm, outer_masks=1.):
  bb = shape[0]
  tt = shape[1]
  pp = shape[2]
  ii = shape[3]
  probs = tf.random_uniform([bb, tt, ii])
  masks = tf.tile(tf.to_float(tf.less(probs, pm))[:, :, None, :], [1, 1, pp, 1])
  return masks * outer_masks
//  [python]
def make_bernoulli_masks(shape, pm, outer_masks=1.):
  bb = shape[0]
  tt = shape[1]
  pp = shape[2]
  ii = shape[3]
  probs = tf.random_uniform([bb, tt, ii])
  masks = tf.tile(tf.cast(tf.less(probs, pm), tf.float32)[:, :, None, :], [1, 1, pp, 1])
  return masks * outer_masks
//  [python]
def apply_cmap(brightness, cmap):
  indices = tf.to_int32(tf.round(brightness * 255))
  cm = matplotlib.cm.get_cmap(cmap)
  colors = tf.constant(cm.colors, dtype=tf.float32)
  return tf.gather(colors, indices)
//  [python]
def apply_cmap(brightness, cmap):
  indices = tf.cast(tf.round(brightness * 255), tf.int32)
  cm = matplotlib.cm.get_cmap(cmap)
  colors = tf.constant(cm.colors, dtype=tf.float32)
  return tf.gather(colors, indices)
//  [python]
def NormalizePixelValues(image,
                         pixel_value_offset=128.0,
                         pixel_value_scale=128.0):
  """Normalize image pixel values.

  Args:
    image: a uint8 tensor.
    pixel_value_offset: a Python float, offset for normalizing pixel values.
    pixel_value_scale: a Python float, scale for normalizing pixel values.

  Returns:
    image: a float32 tensor of the same shape as the input image.
  """
  image = tf.to_float(image)
  image = tf.div(tf.subtract(image, pixel_value_offset), pixel_value_scale)
  return image
//  [python]
def NormalizePixelValues(image,
                         pixel_value_offset=128.0,
                         pixel_value_scale=128.0):
  """Normalize image pixel values.

  Args:
    image: a uint8 tensor.
    pixel_value_offset: a Python float, offset for normalizing pixel values.
    pixel_value_scale: a Python float, scale for normalizing pixel values.

  Returns:
    image: a float32 tensor of the same shape as the input image.
  """
  image = tf.cast(image, tf.float32)
  image = tf.div(tf.subtract(image, pixel_value_offset), pixel_value_scale)
  return image
//  [python]
def CalculateReceptiveBoxes(height, width, rf, stride, padding):
  """Calculate receptive boxes for each feature point.

  Args:
    height: The height of feature map.
    width: The width of feature map.
    rf: The receptive field size.
    stride: The effective stride between two adjacent feature points.
    padding: The effective padding size.
  Returns:
    rf_boxes: [N, 4] receptive boxes tensor. Here N equals to height x width.
    Each box is represented by [ymin, xmin, ymax, xmax].
  """
  x, y = tf.meshgrid(tf.range(width), tf.range(height))
  coordinates = tf.reshape(tf.stack([y, x], axis=2), [-1, 2])
  # [y,x,y,x]
  point_boxes = tf.to_float(tf.concat([coordinates, coordinates], 1))
  bias = [-padding, -padding, -padding + rf - 1, -padding + rf - 1]
  rf_boxes = stride * point_boxes + bias
  return rf_boxes
//  [python]
def CalculateReceptiveBoxes(height, width, rf, stride, padding):
  """Calculate receptive boxes for each feature point.

  Args:
    height: The height of feature map.
    width: The width of feature map.
    rf: The receptive field size.
    stride: The effective stride between two adjacent feature points.
    padding: The effective padding size.
  Returns:
    rf_boxes: [N, 4] receptive boxes tensor. Here N equals to height x width.
    Each box is represented by [ymin, xmin, ymax, xmax].
  """
  x, y = tf.meshgrid(tf.range(width), tf.range(height))
  coordinates = tf.reshape(tf.stack([y, x], axis=2), [-1, 2])
  # [y,x,y,x]
  point_boxes = tf.cast(tf.concat([coordinates, coordinates], 1), tf.float32)
  bias = [-padding, -padding, -padding + rf - 1, -padding + rf - 1]
  rf_boxes = stride * point_boxes + bias
  return rf_boxes
//  [python]
def prep_dataset_fn(filenames, splits, is_training):
  """Prepares training/test datasets in the format required by tf.estimator."""
  if is_training:
    filenames = filenames[splits == 0]
  else:
    filenames = filenames[splits == 1]

  def _load_image(filename):
    """Load an image, center-crop it to a square, and then resize it."""
    image = tf.image.decode_jpeg(tf.read_file(filename), channels=3)
    # Crop to the centermost (IMAGE_PRECROP_SIZE, IMAGE_PRECROP_SIZE) image.
    idx0 = (tf.shape(image)[0] - IMAGE_PRECROP_SIZE) // 2
    idx1 = (tf.shape(image)[1] - IMAGE_PRECROP_SIZE) // 2
    image_cropped = tf.slice(image, [idx0, idx1, 0],
                             [IMAGE_PRECROP_SIZE] * 2 + [-1])
    # Downsample the image to be (IMAGE_SIZE, IMAGE_SIZE).
    image_resized = tf.image.resize_images(image_cropped, [IMAGE_SIZE] * 2)
    image_resized = tf.to_float(image_resized) / 255.
    return image_resized, 0

  def prep_dataset():
    """Return a TF dataset of images expected by tf.estimator."""
    filenames_tf = tf.constant(filenames)
    dataset = tf.data.Dataset.from_tensor_slices(filenames_tf)
    dataset = dataset.map(_load_image)
    if is_training:
      dataset = dataset.shuffle(50000).repeat()
    dataset = dataset.batch(FLAGS.batch_size)
    return dataset

  return prep_dataset
//  [python]
def prep_dataset_fn(filenames, splits, is_training):
  """Prepares training/test datasets in the format required by tf.estimator."""
  if is_training:
    filenames = filenames[splits == 0]
  else:
    filenames = filenames[splits == 1]

  def _load_image(filename):
    """Load an image, center-crop it to a square, and then resize it."""
    image = tf.image.decode_jpeg(tf.read_file(filename), channels=3)
    # Crop to the centermost (IMAGE_PRECROP_SIZE, IMAGE_PRECROP_SIZE) image.
    idx0 = (tf.shape(image)[0] - IMAGE_PRECROP_SIZE) // 2
    idx1 = (tf.shape(image)[1] - IMAGE_PRECROP_SIZE) // 2
    image_cropped = tf.slice(image, [idx0, idx1, 0],
                             [IMAGE_PRECROP_SIZE] * 2 + [-1])
    # Downsample the image to be (IMAGE_SIZE, IMAGE_SIZE).
    image_resized = tf.image.resize_images(image_cropped, [IMAGE_SIZE] * 2)
    image_resized = tf.cast(image_resized, tf.float32) / 255.
    return image_resized, 0

  def prep_dataset():
    """Return a TF dataset of images expected by tf.estimator."""
    filenames_tf = tf.constant(filenames)
    dataset = tf.data.Dataset.from_tensor_slices(filenames_tf)
    dataset = dataset.map(_load_image)
    if is_training:
      dataset = dataset.shuffle(50000).repeat()
    dataset = dataset.batch(FLAGS.batch_size)
    return dataset

  return prep_dataset
//  [python]
def model_fn(features, labels, mode, params):
  """A model is called by TpuEstimator."""
  del labels
  del features

  mesh_shape = mtf.convert_to_shape(FLAGS.mesh_shape)
  layout_rules = mtf.convert_to_layout_rules(FLAGS.layout)

  ctx = params['context']
  num_hosts = ctx.num_hosts
  host_placement_fn = ctx.tpu_host_placement_function
  device_list = [host_placement_fn(host_id=t) for t in range(num_hosts)]
  tf.logging.info('device_list = %s' % device_list,)

  mesh_devices = [''] * mesh_shape.size
  mesh_impl = mtf.simd_mesh_impl.SimdMeshImpl(mesh_shape, layout_rules,
                                              mesh_devices,
                                              ctx.device_assignment)

  graph = mtf.Graph()
  mesh = mtf.Mesh(graph, "fft_mesh")

  with mtf.utils.outside_all_rewrites():
    err = benchmark_model(mesh)

  lowering = mtf.Lowering(graph, {mesh: mesh_impl})

  tf_err = tf.to_float(lowering.export_to_tf_tensor(err))

  with mtf.utils.outside_all_rewrites():
    return tpu_estimator.TPUEstimatorSpec(mode, loss=tf_err)
//  [python]
def model_fn(features, labels, mode, params):
  """A model is called by TpuEstimator."""
  del labels
  del features

  mesh_shape = mtf.convert_to_shape(FLAGS.mesh_shape)
  layout_rules = mtf.convert_to_layout_rules(FLAGS.layout)

  ctx = params['context']
  num_hosts = ctx.num_hosts
  host_placement_fn = ctx.tpu_host_placement_function
  device_list = [host_placement_fn(host_id=t) for t in range(num_hosts)]
  tf.logging.info('device_list = %s' % device_list,)

  mesh_devices = [''] * mesh_shape.size
  mesh_impl = mtf.simd_mesh_impl.SimdMeshImpl(mesh_shape, layout_rules,
                                              mesh_devices,
                                              ctx.device_assignment)

  graph = mtf.Graph()
  mesh = mtf.Mesh(graph, "fft_mesh")

  with mtf.utils.outside_all_rewrites():
    err = benchmark_model(mesh)

  lowering = mtf.Lowering(graph, {mesh: mesh_impl})

  tf_err = tf.cast(lowering.export_to_tf_tensor(err), tf.float32)

  with mtf.utils.outside_all_rewrites():
    return tpu_estimator.TPUEstimatorSpec(mode, loss=tf_err)
//  [python]
def compute_loss_sigmoid(logits, targets, final_sequence_length, target_sequence_length, mode):

    assert mode != tf.estimator.ModeKeys.PREDICT

    if mode == tf.estimator.ModeKeys.TRAIN:
        target_weights = tf.sequence_mask(
            target_sequence_length, dtype=tf.float32)
        loss = sequence_loss_sigmoid(logits, targets, target_weights)
    else:
        max_ts = tf.reduce_max(target_sequence_length)
        max_fs = tf.reduce_max(final_sequence_length)

        max_sequence_length = tf.to_int32(
            tf.maximum(max_ts, max_fs))

        logits = tf.slice(logits, begin=[0, 0, 0], size=[-1, max_fs, -1])

        # pad EOS to make targets and logits have same shape
        targets = tf.pad(targets, [[0, 0], [0, tf.maximum(
            0, max_sequence_length - tf.shape(targets)[1])], [0, 0]], constant_values=0)
        logits = tf.pad(logits, [[0, 0], [0, tf.maximum(
            0, max_sequence_length - tf.shape(logits)[1])], [0, 0]], constant_values=0)

        # find larger length between predictions and targets
        sequence_length = tf.reduce_max(
            [target_sequence_length, final_sequence_length], 0)

        target_weights = tf.sequence_mask(
            sequence_length, maxlen=max_sequence_length, dtype=tf.float32)

        loss = sequence_loss_sigmoid(logits, targets, target_weights)

    return loss
//  [python]
def compute_loss_sigmoid(logits, targets, final_sequence_length, target_sequence_length, mode):

    assert mode != tf.estimator.ModeKeys.PREDICT

    if mode == tf.estimator.ModeKeys.TRAIN:
        target_weights = tf.sequence_mask(
            target_sequence_length, dtype=tf.float32)
        loss = sequence_loss_sigmoid(logits, targets, target_weights)
    else:
        max_ts = tf.reduce_max(target_sequence_length)
        max_fs = tf.reduce_max(final_sequence_length)

        max_sequence_length = tf.cast(
            tf.maximum(max_ts, max_fs), tf.int32)

        logits = tf.slice(logits, begin=[0, 0, 0], size=[-1, max_fs, -1])

        # pad EOS to make targets and logits have same shape
        targets = tf.pad(targets, [[0, 0], [0, tf.maximum(
            0, max_sequence_length - tf.shape(targets)[1])], [0, 0]], constant_values=0)
        logits = tf.pad(logits, [[0, 0], [0, tf.maximum(
            0, max_sequence_length - tf.shape(logits)[1])], [0, 0]], constant_values=0)

        # find larger length between predictions and targets
        sequence_length = tf.reduce_max(
            [target_sequence_length, final_sequence_length], 0)

        target_weights = tf.sequence_mask(
            sequence_length, maxlen=max_sequence_length, dtype=tf.float32)

        loss = sequence_loss_sigmoid(logits, targets, target_weights)

    return loss
//  [python]
def decode_record(record):
  """Decodes a record to a TensorFlow example."""
  name_to_features = {
      "input_ids":
          tf.FixedLenFeature([max_seq_length], tf.int64),
      "input_mask":
          tf.FixedLenFeature([max_seq_length], tf.int64),
      "segment_ids":
          tf.FixedLenFeature([max_seq_length], tf.int64),
      "masked_lm_positions":
          tf.FixedLenFeature([max_predictions_per_seq], tf.int64),
      "masked_lm_ids":
          tf.FixedLenFeature([max_predictions_per_seq], tf.int64),
      "masked_lm_weights":
          tf.FixedLenFeature([max_predictions_per_seq], tf.float32),
      "next_sentence_labels":
          tf.FixedLenFeature([1], tf.int64),
  }

  example = tf.parse_single_example(record, name_to_features)

  # tf.Example only supports tf.int64, but the TPU only supports tf.int32.
  # So cast all int64 to int32.
  for name in list(example.keys()):
    t = example[name]
    if t.dtype == tf.int64:
      t = tf.to_int32(t)
    example[name] = t

  return example
//  [python]
def decode_record(record):
  """Decodes a record to a TensorFlow example."""
  name_to_features = {
      "input_ids":
          tf.FixedLenFeature([max_seq_length], tf.int64),
      "input_mask":
          tf.FixedLenFeature([max_seq_length], tf.int64),
      "segment_ids":
          tf.FixedLenFeature([max_seq_length], tf.int64),
      "masked_lm_positions":
          tf.FixedLenFeature([max_predictions_per_seq], tf.int64),
      "masked_lm_ids":
          tf.FixedLenFeature([max_predictions_per_seq], tf.int64),
      "masked_lm_weights":
          tf.FixedLenFeature([max_predictions_per_seq], tf.float32),
      "next_sentence_labels":
          tf.FixedLenFeature([1], tf.int64),
  }

  example = tf.parse_single_example(record, name_to_features)

  # tf.Example only supports tf.int64, but the TPU only supports tf.int32.
  # So cast all int64 to int32.
  for name in list(example.keys()):
    t = example[name]
    if t.dtype == tf.int64:
      t = tf.cast(t, tf.int32)
    example[name] = t

  return example
//  [python]
def maskrcnn_loss(self, mask_logits, fg_labels, fg_target_masks,
                    scope="maskrcnn_loss"):
    with tf.variable_scope(scope):
      # mask_logits: [N_FG, num_cat, 14, 14]
      # fg_labels: [N_FG]
      # fg_target_masks: [N_FG, 14, 14]
      num_fg = tf.size(fg_labels)
      # [N_FG, 2] # these index is used to get the pos cat"s logit
      indices = tf.stack([tf.range(num_fg), tf.to_int32(fg_labels) - 1], axis=1)
      # ignore other class"s logit
      # [N_FG, 14, 14]
      mask_logits = tf.gather_nd(mask_logits, indices)
      mask_probs = tf.sigmoid(mask_logits)

      loss = tf.nn.sigmoid_cross_entropy_with_logits(
          labels=fg_target_masks, logits=mask_logits)
      loss = tf.reduce_mean(loss, name="maskrcnn_loss")

      return loss
//  [python]
def maskrcnn_loss(self, mask_logits, fg_labels, fg_target_masks,
                    scope="maskrcnn_loss"):
    with tf.variable_scope(scope):
      # mask_logits: [N_FG, num_cat, 14, 14]
      # fg_labels: [N_FG]
      # fg_target_masks: [N_FG, 14, 14]
      num_fg = tf.size(fg_labels)
      # [N_FG, 2] # these index is used to get the pos cat"s logit
      indices = tf.stack([tf.range(num_fg), tf.cast(fg_labels - 1, tf.int32)], axis=1)
      # ignore other class"s logit
      # [N_FG, 14, 14]
      mask_logits = tf.gather_nd(mask_logits, indices)
      mask_probs = tf.sigmoid(mask_logits)

      loss = tf.nn.sigmoid_cross_entropy_with_logits(
          labels=fg_target_masks, logits=mask_logits)
      loss = tf.reduce_mean(loss, name="maskrcnn_loss")

      return loss
//  [python]
def slice_feature_and_anchors(self, image_shape2d, p23456, anchors):
    # anchors is the numpy anchors for different levels
    config = self.config
    # the anchor labels and boxes are grouped into
    gt_anchor_labels = self.anchor_labels
    gt_anchor_boxes = self.anchor_boxes
    self.sliced_anchor_labels = []
    self.sliced_anchor_boxes = []
    for i, stride in enumerate(config.anchor_strides):
      with tf.name_scope("FPN_slice_lvl%s" % (i)):
        if i < 3:
          # Images are padded for p5, which are too large for p2-p4.
          pi = p23456[i]
          target_shape = tf.to_int32(tf.ceil(tf.to_float(image_shape2d) * \
              (1.0 / stride)))

          p23456[i] = tf.slice(
              pi, [0, 0, 0, 0], tf.concat([[-1, -1], target_shape], axis=0))
          p23456[i].set_shape([1, pi.shape[1], None, None])

        shape2d = tf.shape(p23456[i])[2:] # h,W
        slice3d = tf.concat([shape2d, [-1]], axis=0)
        slice4d = tf.concat([shape2d, [-1, -1]], axis=0)

        anchors[i] = tf.slice(anchors[i], [0, 0, 0, 0], slice4d)
        self.sliced_anchor_labels.append(
            tf.slice(gt_anchor_labels[i], [0, 0, 0], slice3d))
        self.sliced_anchor_boxes.append(tf.slice(
            gt_anchor_boxes[i], [0, 0, 0, 0], slice4d))
//  [python]
def slice_feature_and_anchors(self, image_shape2d, p23456, anchors):
    # anchors is the numpy anchors for different levels
    config = self.config
    # the anchor labels and boxes are grouped into
    gt_anchor_labels = self.anchor_labels
    gt_anchor_boxes = self.anchor_boxes
    self.sliced_anchor_labels = []
    self.sliced_anchor_boxes = []
    for i, stride in enumerate(config.anchor_strides):
      with tf.name_scope("FPN_slice_lvl%s" % (i)):
        if i < 3:
          # Images are padded for p5, which are too large for p2-p4.
          pi = p23456[i]
          target_shape = tf.cast(tf.ceil(tf.to_float(image_shape2d) * \
              (1.0 / stride)), tf.int32)

          p23456[i] = tf.slice(
              pi, [0, 0, 0, 0], tf.concat([[-1, -1], target_shape], axis=0))
          p23456[i].set_shape([1, pi.shape[1], None, None])

        shape2d = tf.shape(p23456[i])[2:] # h,W
        slice3d = tf.concat([shape2d, [-1]], axis=0)
        slice4d = tf.concat([shape2d, [-1, -1]], axis=0)

        anchors[i] = tf.slice(anchors[i], [0, 0, 0, 0], slice4d)
        self.sliced_anchor_labels.append(
            tf.slice(gt_anchor_labels[i], [0, 0, 0], slice3d))
        self.sliced_anchor_boxes.append(tf.slice(
            gt_anchor_boxes[i], [0, 0, 0, 0], slice4d))
//  [python]
def fpn_map_rois_to_levels(self, boxes):

    def tf_area(boxes):
      x_min, y_min, x_max, y_max = tf.split(boxes, 4, axis=1)
      return tf.squeeze((y_max - y_min) * (x_max - x_min), [1])

    sqrtarea = tf.sqrt(tf_area(boxes))
    level = tf.to_int32(tf.floor(4 + tf.log(sqrtarea * (1. / 224) + 1e-6) * \
        (1.0 / np.log(2))))
    # RoI levels range from 2~5 (not 6)
    level_ids = [
        tf.where(level <= 2),
        tf.where(tf.equal(level, 3)),# problems with ==?
        tf.where(tf.equal(level, 4)),
        tf.where(level >= 5)]

    level_ids = [tf.reshape(x, [-1], name="roi_level%s_id" % (i + 2))
                 for i, x in enumerate(level_ids)]
    #num_in_levels = [tf.size(x, name="num_roi_level%s" % (i + 2))
    #                 for i, x in enumerate(level_ids)]

    level_boxes = [tf.gather(boxes, ids) for ids in level_ids]
    return level_ids, level_boxes
//  [python]
def fpn_map_rois_to_levels(self, boxes):

    def tf_area(boxes):
      x_min, y_min, x_max, y_max = tf.split(boxes, 4, axis=1)
      return tf.squeeze((y_max - y_min) * (x_max - x_min), [1])

    sqrtarea = tf.sqrt(tf_area(boxes))
    level = tf.cast(tf.floor(4 + tf.log(sqrtarea * (1. / 224) + 1e-6) * \
        (1.0 / np.log(2))), tf.int32)
    # RoI levels range from 2~5 (not 6)
    level_ids = [
        tf.where(level <= 2),
        tf.where(tf.equal(level, 3)),# problems with ==?
        tf.where(tf.equal(level, 4)),
        tf.where(level >= 5)]

    level_ids = [tf.reshape(x, [-1], name="roi_level%s_id" % (i + 2))
                 for i, x in enumerate(level_ids)]
    #num_in_levels = [tf.size(x, name="num_roi_level%s" % (i + 2))
    #                 for i, x in enumerate(level_ids)]

    level_boxes = [tf.gather(boxes, ids) for ids in level_ids]
    return level_ids, level_boxes
//  [python]
def fpn_map_rois_to_levels(self, boxes):

    # boxes: [K, 5] -> each level got [M, 5]
    # [M]
    sqrtarea = tf.sqrt(self.tf_area(boxes))
    level = tf.to_int32(tf.floor(4 + tf.log(sqrtarea * (1. / 224) + 1e-6) * \
        (1.0 / np.log(2))))

    # RoI levels range from 2~5 (not 6)\
    # [M, 1] for [K] shape level
    # [M, 2] for [B, K] level
    level_idxs = [
        tf.where(level <= 2),
        tf.where(tf.equal(level, 3)),
        tf.where(tf.equal(level, 4)),
        tf.where(level >= 5)]

    level_idxs = [tf.reshape(x, [-1], name="roi_level%s_id" % (i + 2))
                  for i, x in enumerate(level_idxs)]

    # for each level, get [M*, 5] boxes
    level_boxes = [tf.gather(boxes, ids) for ids in level_idxs]
    return level_idxs, level_boxes
//  [python]
def fpn_map_rois_to_levels(self, boxes):

    # boxes: [K, 5] -> each level got [M, 5]
    # [M]
    sqrtarea = tf.sqrt(self.tf_area(boxes))
    level = tf.cast(tf.floor(4 + tf.log(sqrtarea * (1. / 224) + 1e-6) * \
        (1.0 / np.log(2))), tf.int32)

    # RoI levels range from 2~5 (not 6)\
    # [M, 1] for [K] shape level
    # [M, 2] for [B, K] level
    level_idxs = [
        tf.where(level <= 2),
        tf.where(tf.equal(level, 3)),
        tf.where(tf.equal(level, 4)),
        tf.where(level >= 5)]

    level_idxs = [tf.reshape(x, [-1], name="roi_level%s_id" % (i + 2))
                  for i, x in enumerate(level_idxs)]

    # for each level, get [M*, 5] boxes
    level_boxes = [tf.gather(boxes, ids) for ids in level_idxs]
    return level_idxs, level_boxes
//  [python]
def slice_feature_and_anchors(self, image_shape2d, p23456, anchors):
    # p2345 each is [B, C, FS, FS]
    # anchors is the numpy anchors for different levels
    config = self.config
    batch_size = p23456[0].get_shape()[0]
    # the anchor labels and boxes are grouped into
    gt_anchor_labels = self.anchor_labels
    gt_anchor_boxes = self.anchor_boxes
    self.sliced_anchor_labels = []
    self.sliced_anchor_boxes = []
    for i, stride in enumerate(config.anchor_strides):
      with tf.name_scope("FPN_slice_lvl%s" % (i)):
        if i < 3:
          # Images are padded for p5, which are too large for p2-p4.
          pi = p23456[i]
          target_shape = tf.to_int32(tf.ceil(tf.to_float(image_shape2d) * \
              (1.0 / stride)))

          p23456[i] = tf.slice(
              pi, [0, 0, 0, 0], tf.concat([[-1, -1], target_shape], axis=0))
          p23456[i].set_shape([batch_size, pi.shape[1], None, None])

        shape2d = tf.shape(p23456[i])[2:] # h,W
        slice3d = tf.concat([shape2d, [-1]], axis=0)
        slice4d = tf.concat([shape2d, [-1, -1]], axis=0)

        anchors[i] = tf.slice(anchors[i], [0, 0, 0, 0], slice4d)

        self.sliced_anchor_labels.append(
            tf.slice(gt_anchor_labels[i], [0, 0, 0], slice3d))
        self.sliced_anchor_boxes.append(tf.slice(
            gt_anchor_boxes[i], [0, 0, 0, 0], slice4d))
//  [python]
def slice_feature_and_anchors(self, image_shape2d, p23456, anchors):
    # p2345 each is [B, C, FS, FS]
    # anchors is the numpy anchors for different levels
    config = self.config
    batch_size = p23456[0].get_shape()[0]
    # the anchor labels and boxes are grouped into
    gt_anchor_labels = self.anchor_labels
    gt_anchor_boxes = self.anchor_boxes
    self.sliced_anchor_labels = []
    self.sliced_anchor_boxes = []
    for i, stride in enumerate(config.anchor_strides):
      with tf.name_scope("FPN_slice_lvl%s" % (i)):
        if i < 3:
          # Images are padded for p5, which are too large for p2-p4.
          pi = p23456[i]
          target_shape = tf.cast(tf.ceil(tf.to_float(image_shape2d) * \
              (1.0 / stride)), tf.int32)

          p23456[i] = tf.slice(
              pi, [0, 0, 0, 0], tf.concat([[-1, -1], target_shape], axis=0))
          p23456[i].set_shape([batch_size, pi.shape[1], None, None])

        shape2d = tf.shape(p23456[i])[2:] # h,W
        slice3d = tf.concat([shape2d, [-1]], axis=0)
        slice4d = tf.concat([shape2d, [-1, -1]], axis=0)

        anchors[i] = tf.slice(anchors[i], [0, 0, 0, 0], slice4d)

        self.sliced_anchor_labels.append(
            tf.slice(gt_anchor_labels[i], [0, 0, 0], slice3d))
        self.sliced_anchor_boxes.append(tf.slice(
            gt_anchor_boxes[i], [0, 0, 0, 0], slice4d))
//  [python]
def setup_loss(agent_net, targets_tensor):
    l2_norm = tf.global_norm(tf.trainable_variables())
    # steering_error = tf.reduce_mean(tf.abs(agent_net.out[:, 4] - targets_tensor[:, 4]))
    # throttle_error = tf.reduce_mean(tf.abs(agent_net.out[:, 5] - targets_tensor[:, 5]))
    # tf.summary.scalar("steering_error/train", steering_error)
    # tf.summary.scalar("throttle_error/train", throttle_error)
    # l2_norm = tf.Print(l2_norm, [steering_error], 'TRAIN STEERING ERROR IS!!!!!!!!! ', summarize=100)
    loss = 0.5 * tf.reduce_sum(tf.square(agent_net.out - targets_tensor)) / tf.to_float(tf.shape(agent_net.input)[0])
    tf.summary.scalar("model/loss", loss)
    tf.summary.scalar("model/l2_norm", l2_norm)
    total_loss = loss + agent_net.weight_decay * l2_norm
    tf.summary.scalar("model/total_loss", total_loss)
    return total_loss
//  [python]
def setup_loss(agent_net, targets_tensor):
    l2_norm = tf.global_norm(tf.trainable_variables())
    # steering_error = tf.reduce_mean(tf.abs(agent_net.out[:, 4] - targets_tensor[:, 4]))
    # throttle_error = tf.reduce_mean(tf.abs(agent_net.out[:, 5] - targets_tensor[:, 5]))
    # tf.summary.scalar("steering_error/train", steering_error)
    # tf.summary.scalar("throttle_error/train", throttle_error)
    # l2_norm = tf.Print(l2_norm, [steering_error], 'TRAIN STEERING ERROR IS!!!!!!!!! ', summarize=100)
    loss = 0.5 * tf.reduce_sum(tf.square(agent_net.out - targets_tensor)) / tf.cast(tf.shape(agent_net.input)[0], tf.float32)
    tf.summary.scalar("model/loss", loss)
    tf.summary.scalar("model/l2_norm", l2_norm)
    total_loss = loss + agent_net.weight_decay * l2_norm
    tf.summary.scalar("model/total_loss", total_loss)
    return total_loss