jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def binary_classifier(features, labels, mode, params):
    output, label_vocabulary = create_output(features, params, mode)
    predicted_classes = tf.cast(output > 0.5, tf.int64)

    if mode == tf.estimator.ModeKeys.PREDICT:
        predictions = {
            'class_ids': predicted_classes[:, tf.newaxis],
            'probabilities': tf.nn.sigmoid(output) if 'sigmoid' not in output.name.lower() else output,
            'logits': output,
            'classes': get_label_classes(predicted_classes, label_vocabulary)
        }
        return tf.estimator.EstimatorSpec(mode, predictions=predictions)

    label_ids = get_label_ids(labels, label_vocabulary)
    loss = getattr(tensorflow.losses, params['loss_function'])(tf.reshape(label_ids, tf.shape(output)), output)

    accuracy = tf.metrics.accuracy(labels=label_ids,
                                   predictions=predicted_classes,
                                   name='accuracy')
    probs = tf.nn.sigmoid(output) if 'sigmoid' not in output.name.lower() else output
    auc = tf.metrics.auc(labels=label_ids,
                         predictions=probs,
                         name='auc')
    auc_pr = tf.metrics.auc(labels=label_ids,
                            predictions=probs,
                            name='auc_precision_recall',
                            curve='PR')
    precision = tf.metrics.precision(labels=label_ids,
                                     predictions=probs,
                                     name='precision')
    recall = tf.metrics.recall(labels=label_ids,
                               predictions=probs,
                               name='recall')
    metrics = {
        'accuracy': accuracy,
        'auc': auc,
        'auc_precision_recall': auc_pr,
        'precision': precision,
        'recall': recall
    }
    tf.summary.scalar('accuracy', accuracy[1])
    tf.summary.scalar('auc', auc[1])
    tf.summary.scalar('auc_precision_recall', auc_pr[1])
    tf.summary.scalar('precision', precision[1])
    tf.summary.scalar('recall', recall[1])

    if mode == tf.estimator.ModeKeys.EVAL:
        return tf.estimator.EstimatorSpec(mode, loss=loss, eval_metric_ops=metrics)

    train_op = create_train_op(loss, params, mode)

    return tf.estimator.EstimatorSpec(mode, loss=loss, train_op=train_op)
//  [python]
def padded_one_hot_encoding(indices, depth, left_pad):
  """Returns a zero padded one-hot tensor.

  This function converts a sparse representation of indices (e.g., [4]) to a
  zero padded one-hot representation (e.g., [0, 0, 0, 0, 1] with depth = 4 and
  left_pad = 1). If `indices` is empty, the result will simply be a tensor of
  shape (0, depth + left_pad). If depth = 0, then this function just returns
  `None`.

  Args:
    indices: an integer tensor of shape [num_indices].
    depth: depth for the one-hot tensor (integer).
    left_pad: number of zeros to left pad the one-hot tensor with (integer).

  Returns:
    padded_onehot: a tensor with shape (num_indices, depth + left_pad). Returns
      `None` if the depth is zero.

  Raises:
    ValueError: if `indices` does not have rank 1 or if `left_pad` or `depth are
      either negative or non-integers.

  TODO: add runtime checks for depth and indices.
  """
  if depth < 0 or not isinstance(depth, six.integer_types):
    raise ValueError('`depth` must be a non-negative integer.')
  if left_pad < 0 or not isinstance(left_pad, six.integer_types):
    raise ValueError('`left_pad` must be a non-negative integer.')
  if depth == 0:
    return None
  if len(indices.get_shape().as_list()) != 1:
    raise ValueError('`indices` must have rank 1')

  def one_hot_and_pad():
    one_hot = tf.cast(tf.one_hot(tf.cast(indices, tf.int64), depth,
                                 on_value=1, off_value=0), tf.float32)
    return tf.pad(one_hot, [[0, 0], [left_pad, 0]], mode='CONSTANT')
  result = tf.cond(tf.greater(tf.size(indices), 0), one_hot_and_pad,
                   lambda: tf.zeros((depth + left_pad, 0)))
  return tf.reshape(result, [-1, depth + left_pad])
//  [python]
def retrieve_seq_length_op3(data, pad_val=0): # HangSheng: return tensor for sequence length, if input is tf.string
    data_shape_size = data.get_shape().ndims
    if data_shape_size == 3:
        return tf.reduce_sum(tf.to_int32(tf.reduce_any(tf.not_equal(data, pad_val), axis=2)), 1)
    elif data_shape_size == 2:
        return tf.reduce_sum(tf.to_int32(tf.not_equal(data, pad_val)), 1)
    elif data_shape_size == 1:
        raise ValueError("retrieve_seq_length_op3: data has wrong shape!")
    else:
        raise ValueError("retrieve_seq_length_op3: handling data_shape_size %s hasn't been implemented!" % (data_shape_size))
//  [python]
def retrieve_seq_length_op3(data, pad_val=0): # HangSheng: return tensor for sequence length, if input is tf.string
    data_shape_size = data.get_shape().ndims
    if data_shape_size == 3:
        return tf.reduce_sum(tf.cast(tf.reduce_any(tf.not_equal(data, pad_val), axis=2), dtype=tf.int32), 1)
    elif data_shape_size == 2:
        return tf.reduce_sum(tf.cast(tf.not_equal(data, pad_val), dtype=tf.int32), 1)
    elif data_shape_size == 1:
        raise ValueError("retrieve_seq_length_op3: data has wrong shape!")
    else:
        raise ValueError("retrieve_seq_length_op3: handling data_shape_size %s hasn't been implemented!" % (data_shape_size))
//  [python]
def get_OHEM_mask(loss_tensor, selected_num):
    shape = loss_tensor.get_shape().as_list()
    ce_loss_tensor_flatten = tf.reshape(loss_tensor, [shape[0] * shape[1]])
    selected_values, selected_indices = tf.nn.top_k(ce_loss_tensor_flatten, k=selected_num)
    selected_mask = tf.to_float(tf.greater_equal(loss_tensor, selected_values[-1]))
    return selected_mask, selected_indices
//  [python]
def get_OHEM_mask(loss_tensor, selected_num):
    shape = loss_tensor.get_shape().as_list()
    ce_loss_tensor_flatten = tf.reshape(loss_tensor, [shape[0] * shape[1]])
    selected_values, selected_indices = tf.nn.top_k(ce_loss_tensor_flatten, k=selected_num)
    selected_mask = tf.cast(tf.greater_equal(loss_tensor, selected_values[-1]), tf.float32)
    return selected_mask, selected_indices
//  [python]
def target_mask_op(data, pad_val=0):        # HangSheng: return tensor for mask,if input is tf.string
    data_shape_size = data.get_shape().ndims
    if data_shape_size == 3:
        return tf.to_int32(tf.reduce_any(tf.not_equal(data, pad_val), axis=2))
    elif data_shape_size == 2:
        return tf.to_int32(tf.not_equal(data, pad_val))
    elif data_shape_size == 1:
        raise ValueError("target_mask_op: data has wrong shape!")
    else:
        raise ValueError("target_mask_op: handling data_shape_size %s hasn't been implemented!" % (data_shape_size))
//  [python]
def target_mask_op(data, pad_val=0):        # HangSheng: return tensor for mask,if input is tf.string
    data_shape_size = data.get_shape().ndims
    if data_shape_size == 3:
        return tf.cast(tf.reduce_any(tf.not_equal(data, pad_val), axis=2), dtype=tf.int32)
    elif data_shape_size == 2:
        return tf.cast(tf.not_equal(data, pad_val), dtype=tf.int32)
    elif data_shape_size == 1:
        raise ValueError("target_mask_op: data has wrong shape!")
    else:
        raise ValueError("target_mask_op: handling data_shape_size %s hasn't been implemented!" % (data_shape_size))
//  [python]
def build_loss(self, mask_tensor, cross_entropy_coff=1.0, dice_coff=1.0, focal_loss_coff=1.0, triplet_loss_coff=1.0):
        print('start build loss')
        if self.prediction is None:
            return None
        ce_loss = tf.keras.losses.sparse_categorical_crossentropy(mask_tensor, self.prediction)
        tf.summary.scalar('loss/final_ce', tf.reduce_mean(ce_loss))

        dice_loss = dice_coef_loss(tf.to_float(mask_tensor), self.prediction[:, :, :, 1])
        print('dice_loss is ', dice_loss)
        tf.summary.scalar('loss/final_dice_loss', dice_loss)

        focal_loss_t = focal_loss()(tf.to_float(mask_tensor), self.prediction[:, :, :, 1])
        tf.summary.scalar('loss/focal_loss', tf.reduce_mean(focal_loss_t))

        # triplet_loss_t = tf.keras.layers.Lambda(lambda x: triplet_loss_OHEM(*x))((self.prediction, mask_tensor,
        #                                                                           self.decoder_outputs[-2]))
        # print('triplet_loss_t is ', triplet_loss_t)
        # tf.summary.scalar('loss/triplet_loss', triplet_loss_t)

        final_loss = cross_entropy_coff * tf.reduce_mean(
            ce_loss) + dice_coff * dice_loss + focal_loss_coff * tf.reduce_mean(
            focal_loss_t)

        return final_loss
//  [python]
def build_loss(self, mask_tensor, cross_entropy_coff=1.0, dice_coff=1.0, focal_loss_coff=1.0, triplet_loss_coff=1.0):
        print('start build loss')
        if self.prediction is None:
            return None
        ce_loss = tf.keras.losses.sparse_categorical_crossentropy(mask_tensor, self.prediction)
        tf.summary.scalar('loss/final_ce', tf.reduce_mean(ce_loss))

        dice_loss = dice_coef_loss(tf.cast(mask_tensor, tf.float32), self.prediction[:, :, :, 1])
        print('dice_loss is ', dice_loss)
        tf.summary.scalar('loss/final_dice_loss', dice_loss)

        focal_loss_t = focal_loss()(tf.cast(mask_tensor, tf.float32), self.prediction[:, :, :, 1])
        tf.summary.scalar('loss/focal_loss', tf.reduce_mean(focal_loss_t))

        # triplet_loss_t = tf.keras.layers.Lambda(lambda x: triplet_loss_OHEM(*x))((self.prediction, mask_tensor,
        #                                                                           self.decoder_outputs[-2]))
        # print('triplet_loss_t is ', triplet_loss_t)
        # tf.summary.scalar('loss/triplet_loss', triplet_loss_t)

        final_loss = cross_entropy_coff * tf.reduce_mean(
            ce_loss) + dice_coff * dice_loss + focal_loss_coff * tf.reduce_mean(
            focal_loss_t)

        return final_loss
//  [python]
def call(self, inputs: tf.Tensor, inputs_positions: Any = None) -> tf.Tensor:
    # inputs.shape is (batch_size, time_len, seq_len, emb_dim).
    del inputs_positions
    pos_embedding_time = self._pos_embedding_time
    if inputs.shape[1] != pos_embedding_time.shape[1]:
      pos_embedding_time = self._interpolate(
          pos_embedding_time,
          from_shape=(1, self.posemb_origin_shape[0]),
          to_shape=(1, self.posemb_target_shape[0]))

    pos_embedding_space = self._pos_embedding_space
    if inputs.shape[2] != pos_embedding_space.shape[1]:
      pos_embedding_space = self._interpolate(
          pos_embedding_space,
          from_shape=(1, self.posemb_origin_shape[1]),
          to_shape=(1, self.posemb_target_shape[1]))

    pos_embedding_time = tf.to_float(pos_embedding_time[:, :, None, :])
    pos_embedding_space = tf.to_float(pos_embedding_space[:, None, :, :])
    return inputs + pos_embedding_time + pos_embedding_space
//  [python]
def call(self, inputs: tf.Tensor, inputs_positions: Any = None) -> tf.Tensor:
    # inputs.shape is (batch_size, time_len, seq_len, emb_dim).
    del inputs_positions
    pos_embedding_time = self._pos_embedding_time
    if inputs.shape[1] != pos_embedding_time.shape[1]:
      pos_embedding_time = self._interpolate(
          pos_embedding_time,
          from_shape=(1, self.posemb_origin_shape[0]),
          to_shape=(1, self.posemb_target_shape[0]))

    pos_embedding_space = self._pos_embedding_space
    if inputs.shape[2] != pos_embedding_space.shape[1]:
      pos_embedding_space = self._interpolate(
          pos_embedding_space,
          from_shape=(1, self.posemb_origin_shape[1]),
          to_shape=(1, self.posemb_target_shape[1]))

    pos_embedding_time = tf.cast(pos_embedding_time[:, :, None, :],
                                 inputs.dtype)
    pos_embedding_space = tf.cast(pos_embedding_space[:, None, :, :],
                                  inputs.dtype)
    return inputs + pos_embedding_time + pos_embedding_space
//  [python]
def build_loss(self, mask_tensor, gen_tensor, cross_entropy_coff=1.0, dice_coff=1.0, focal_loss_coff=1.0):
        print('start build loss')
        if self.prediction is None:
            return None
        ce_loss = tf.keras.losses.sparse_categorical_crossentropy(mask_tensor, self.prediction)
        tf.summary.scalar('loss/sg/final_ce', tf.reduce_mean(ce_loss))

        dice_loss = dice_coef_loss(tf.to_float(mask_tensor), self.prediction[:, :, :, 1])
        print('dice_loss is ', dice_loss)
        tf.summary.scalar('loss/sg/final_dice_loss', dice_loss)

        focal_loss_t = focal_loss()(tf.to_float(mask_tensor), self.prediction[:, :, :, 1])
        tf.summary.scalar('loss/sg/focal_loss', tf.reduce_mean(focal_loss_t))
        final_sg_loss = cross_entropy_coff * tf.reduce_mean(
            ce_loss) + dice_coff * dice_loss + focal_loss_coff * tf.reduce_mean(focal_loss_t)

        cls_ce_loss = tf.keras.losses.sparse_categorical_crossentropy(gen_tensor, self.cls_prediction)
        tf.summary.scalar('loss/cls/final_ce', tf.reduce_mean(cls_ce_loss))

        final_loss = final_sg_loss + tf.reduce_mean(cls_ce_loss)
        return final_loss
//  [python]
def build_loss(self, mask_tensor, gen_tensor, cross_entropy_coff=1.0, dice_coff=1.0, focal_loss_coff=1.0):
        print('start build loss')
        if self.prediction is None:
            return None
        ce_loss = tf.keras.losses.sparse_categorical_crossentropy(mask_tensor, self.prediction)
        tf.summary.scalar('loss/sg/final_ce', tf.reduce_mean(ce_loss))

        dice_loss = dice_coef_loss(tf.cast(mask_tensor, tf.float32), self.prediction[:, :, :, 1])
        print('dice_loss is ', dice_loss)
        tf.summary.scalar('loss/sg/final_dice_loss', dice_loss)

        focal_loss_t = focal_loss()(tf.cast(mask_tensor, tf.float32), self.prediction[:, :, :, 1])
        tf.summary.scalar('loss/sg/focal_loss', tf.reduce_mean(focal_loss_t))
        final_sg_loss = cross_entropy_coff * tf.reduce_mean(
            ce_loss) + dice_coff * dice_loss + focal_loss_coff * tf.reduce_mean(focal_loss_t)

        cls_ce_loss = tf.keras.losses.sparse_categorical_crossentropy(gen_tensor, self.cls_prediction)
        tf.summary.scalar('loss/cls/final_ce', tf.reduce_mean(cls_ce_loss))

        final_loss = final_sg_loss + tf.reduce_mean(cls_ce_loss)
        return final_loss
//  [python]
def call(self, inputs: tf.Tensor, training: bool = None):  # pytype: disable=signature-mismatch  # overriding-parameter-count-checks
    """Forward pass mask scoring head.

    Args:
      inputs: A `tf.Tensor` of the shape [batch_size, width, size, num_classes],
      representing the segmentation logits.
      training: a `bool` indicating whether it is in `training` mode.

    Returns:
      mask_scores: A `tf.Tensor` of predicted mask scores
        [batch_size, num_classes].
    """
    x = tf.stop_gradient(inputs)
    for conv, bn in zip(self._convs, self._conv_norms):
      x = conv(x)
      x = bn(x)
      x = self._activation(x)

    # Casts feat to float32 so the resize op can be run on TPU.
    x = tf.to_float(x)
    x = tf.image.resize(x, size=self._config_dict['fc_input_size'],
                        method=tf.image.ResizeMethod.BILINEAR)
    # Casts it back to be compatible with the rest opetations.
    x = tf.cast(x, inputs.dtype)

    _, h, w, filters = x.get_shape().as_list()
    x = tf.reshape(x, [-1, h * w * filters])

    for fc, bn in zip(self._fcs, self._fc_norms):
      x = fc(x)
      x = bn(x)
      x = self._activation(x)

    ious = self._classifier(x)
    return ious
//  [python]
def call(self, inputs: tf.Tensor, training: bool = None):  # pytype: disable=signature-mismatch  # overriding-parameter-count-checks
    """Forward pass mask scoring head.

    Args:
      inputs: A `tf.Tensor` of the shape [batch_size, width, size, num_classes],
      representing the segmentation logits.
      training: a `bool` indicating whether it is in `training` mode.

    Returns:
      mask_scores: A `tf.Tensor` of predicted mask scores
        [batch_size, num_classes].
    """
    x = tf.stop_gradient(inputs)
    for conv, bn in zip(self._convs, self._conv_norms):
      x = conv(x)
      x = bn(x)
      x = self._activation(x)

    # Casts feat to float32 so the resize op can be run on TPU.
    x = tf.cast(x, tf.float32)
    x = tf.image.resize(x, size=self._config_dict['fc_input_size'],
                        method=tf.image.ResizeMethod.BILINEAR)
    # Casts it back to be compatible with the rest opetations.
    x = tf.cast(x, inputs.dtype)

    _, h, w, filters = x.get_shape().as_list()
    x = tf.reshape(x, [-1, h * w * filters])

    for fc, bn in zip(self._fcs, self._fc_norms):
      x = fc(x)
      x = bn(x)
      x = self._activation(x)

    ious = self._classifier(x)
    return ious
//  [python]
def _build_clk_model(self):
        with self._make_scope('clk_model', self._bf16):
            if self._bf16:
                self._input_layer = tf.cast(self._input_layer, dtype=tf.bfloat16)
            d_clk = self._input_layer
            for layer_id, num_hidden_units in enumerate(self._mlp):
                d_clk = self._create_dense_layer(d_clk, num_hidden_units, tf.nn.relu, f'd{layer_id}_clk')

            d_clk = self._create_dense_layer(d_clk, 1, None, 'output_clk')
            if self._bf16:
                d_clk = tf.cast(d_clk, tf.float32)

        Y_clk = tf.squeeze(d_clk)

        return tf.math.sigmoid(d_clk), Y_clk
//  [python]
def _build_buy_model(self):
        with self._make_scope('buy_model', self._bf16):
            if self._bf16:
                self._input_layer = tf.cast(self._input_layer, dtype=tf.bfloat16)
            d_buy = self._input_layer
            for layer_id, num_hidden_units in enumerate(self._mlp):
                d_buy = self._create_dense_layer(d_buy, num_hidden_units, tf.nn.relu, f'd{layer_id}_buy')

            d_buy = self._create_dense_layer(d_buy, 1, None, 'output_buy')
            if self._bf16:
                d_buy = tf.cast(d_buy, tf.float32)

        Y_buy = tf.squeeze(d_buy)

        return tf.math.sigmoid(d_buy), Y_buy
//  [python]
def _create_model(self):
        with tf.variable_scope('input_layer',
                                partitioner=self._input_layer_partitioner,
                                reuse=tf.AUTO_REUSE):
            if not self._tf and self._adaptive_emb:
                '''Adaptive Embedding Feature Part 1 of 2'''
                adaptive_mask_tensors = {}
                for i in range(len(HASH_INPUTS)):
                    adaptive_mask_tensors[HASH_INPUTS[i]] = tf.ones([self._batch_size],
                                                                    tf.int32)
                self._input_layer = tf.feature_column.input_layer(
                        self.feature,
                        self._feature_columns,
                        adaptive_mask_tensors=adaptive_mask_tensors)
            else:
                self._input_layer = tf.feature_column.input_layer(self.feature, self._feature_columns)

            if self._bf16:
                self._input_layer = tf.cast(self._input_layer, dtype=tf.bfloat16)

        with tf.variable_scope('SimpleMultiTask',
                               reuse=tf.AUTO_REUSE):
            self.clk_model, self.clk_logits = self._build_clk_model()
            self.buy_model, self.buy_logits = self._build_buy_model()
            self._logits = tf.squeeze(tf.stack([self.clk_model, self.buy_model], axis=1), [-1])
            self.probability = self._logits
            self.output = tf.round(self.probability)
//  [python]
def log_probs(self, states, actions, state_reprs, contexts=None):
    assert contexts is not None
    batch_dims = [tf.shape(states)[0], tf.shape(states)[1]]
    contexts = self.tf_context.context_multi_transition_fn(
        contexts, states=t.to_float(state_reprs))

    flat_states = tf.reshape(states,
                             [batch_dims[0] * batch_dims[1], states.shape[-1]])
    flat_contexts = [tf.reshape(t.cast(context, states.dtype),
                                [batch_dims[0] * batch_dims[1], context.shape[-1]])
                     for context in contexts]
    flat_pred_actions = self.actions(flat_states, flat_contexts)
    pred_actions = tf.reshape(flat_pred_actions,
                              batch_dims + [flat_pred_actions.shape[-1]])

    error = tf.square(actions - pred_actions)
    spec_range = (self._action_spec.maximum - self._action_spec.minimum) / 2
    normalized_error = t.cast(error, tf.float64) / tf.constant(spec_range) ** 2
    return -normalized_error
//  [python]
def log_probs(self, states, actions, state_reprs, contexts=None):
    assert contexts is not None
    batch_dims = [tf.shape(states)[0], tf.shape(states)[1]]
    contexts = self.tf_context.context_multi_transition_fn(
        contexts, states=tf.to_float(state_reprs))

    flat_states = tf.reshape(states,
                             [batch_dims[0] * batch_dims[1], states.shape[-1]])
    flat_contexts = [tf.reshape(tf.cast(context, states.dtype),
                                [batch_dims[0] * batch_dims[1], context.shape[-1]])
                     for context in contexts]
    flat_pred_actions = self.actions(flat_states, flat_contexts)
    pred_actions = tf.reshape(flat_pred_actions,
                              batch_dims + [flat_pred_actions.shape[-1]])

    error = tf.square(actions - pred_actions)
    spec_range = (self._action_spec.maximum - self._action_spec.minimum) / 2
    normalized_error = tf.cast(error, tf.float64) / tf.constant(spec_range) ** 2
    return -normalized_error
//  [python]
def sample(self, states, next_states, num_samples, orig_goals, sc=0.5):
    goal_dim = orig_goals.shape[-1]
    spec_range = (self._spec.maximum - self._spec.minimum) / 2 * t.ones([goal_dim])
    loc = t.cast(next_states - states, t.float32)[:, :goal_dim]
    scale = sc * t.tile(t.reshape(spec_range, [1, goal_dim]),
                         [t.shape(states)[0], 1])
    dist = t.distributions.Normal(loc, scale)
    if num_samples == 1:
      return dist.sample()
    samples = t.concat([dist.sample(num_samples - 2),
                         t.expand_dims(loc, 0),
                         t.expand_dims(orig_goals, 0)], 0)
    return uvf_utils.clip_to_spec(samples, self._spec)
//  [python]
def sample(self, states, next_states, num_samples, orig_goals, sc=0.5):
    goal_dim = orig_goals.shape[-1]
    spec_range = (self._spec.maximum - self._spec.minimum) / 2 * tf.ones([goal_dim])
    loc = tf.cast(next_states - states, tf.float32)[:, :goal_dim]
    scale = sc * tf.tile(tf.reshape(spec_range, [1, goal_dim]),
                         [tf.shape(states)[0], 1])
    dist = tf.distributions.Normal(loc, scale)
    if num_samples == 1:
      return dist.sample()
    samples = tf.concat([dist.sample(num_samples - 2),
                         tf.expand_dims(loc, 0),
                         tf.expand_dims(orig_goals, 0)], 0)
    return uvf_utils.clip_to_spec(samples, self._spec)
//  [python]
def pathlen(path):
    with tf.name_scope('path_len'):
        populated = tf.sign(tf.reduce_sum(tf.abs(path), 2))
        length = tf.cast(tf.reduce_sum(populated, 1), tf.int32)
        mask = tf.cast(populated, tf.float32)

    return length, mask
//  [python]
def seqlen(sequence):
    """
    Get true length of sequences (without padding), and mask for true-length in max-length.
    Input of shape: (batch_size, max_seq_length, hidden_dim)
    Output shapes, 
    length: (batch_size)
    mask: (batch_size, max_seq_length, 1)
    """
    with tf.name_scope('seq_len'):
        populated = tf.sign(tf.abs(sequence))
        length = tf.to_int32(tf.reduce_sum(populated, axis=1))
        mask = tf.to_float(tf.expand_dims(populated, -1))

    return length, mask
//  [python]
def seqlen(sequence):
    """
    Get true length of sequences (without padding), and mask for true-length in max-length.
    Input of shape: (batch_size, max_seq_length, hidden_dim)
    Output shapes, 
    length: (batch_size)
    mask: (batch_size, max_seq_length, 1)
    """
    with tf.name_scope('seq_len'):
        populated = tf.sign(tf.abs(sequence))
        length = tf.cast(tf.reduce_sum(populated, axis=1), tf.int32)
        mask = tf.cast(tf.expand_dims(populated, -1), tf.float32)

    return length, mask
//  [python]
def residuals(self, y_true, y_pred):
        y_pred = tf.convert_to_tensor(y_pred)
        y_true = tf.to_float(y_true)
        return y_true - y_pred
//  [python]
def residuals(self, y_true, y_pred):
        y_pred = tf.convert_to_tensor(y_pred)
        y_true = tf.cast(y_true, y_pred.dtype)
        return y_true - y_pred
//  [python]
def get_box_in_a_single_layer(self,datas,threshold):
        bboxes,clses,scores,_ = self.box2box_transform.apply_deltas(datas)
        mask = tf.to_int32(tf.greater_equal(scores,threshold))
        length = tf.reduce_sum(mask,axis=-1)
        return bboxes,clses,scores,length
//  [python]
def get_box_in_a_single_layer(self,datas,threshold):
        bboxes,clses,scores,_ = self.box2box_transform.apply_deltas(datas)
        mask = tf.cast(tf.greater_equal(scores,threshold),tf.int32)
        length = tf.reduce_sum(mask,axis=-1)
        return bboxes,clses,scores,length
//  [python]
def call(self, y_true, y_pred):
        y_pred = tf.convert_to_tensor(y_pred)
        y_true = tf.cast(y_true, y_pred.dtype)
        sq_diff = tf.math.squared_difference(y_true, y_pred)
        return tf.math.reduce_mean(sq_diff, axis=1)
//  [python]
def loss(self, states, next_states, low_actions, low_states):
    batch_size = tf.shape(states)[0]
    d = int(low_states.shape[1])
    # Sample indices into meta-transition to train on.
    probs = 0.99 ** tf.range(d, dtype=tf.float32)
    probs *= tf.constant([1.0] * (d - 1) + [1.0 / (1 - 0.99)],
                         dtype=tf.float32)
    probs /= tf.reduce_sum(probs)
    index_dist = tf.distributions.Categorical(probs=probs, dtype=tf.int64)
    indices = index_dist.sample(batch_size)
    batch_size = tf.to_int64(batch_size)
    next_indices = tf.concat(
        [tf.range(batch_size, dtype=tf.int64)[:, None],
         (1 + indices[:, None]) % d], -1)
    new_next_states = tf.where(indices < d - 1,
                               tf.gather_nd(low_states, next_indices),
                               next_states)
    next_states = new_next_states

    embed1 = tf.to_float(self._state_preprocess_net(states))
    embed2 = tf.to_float(self._state_preprocess_net(next_states))
    action_embed = self._action_embed_net(
        tf.layers.flatten(low_actions), states=states)

    tau = 2.0
    fn = lambda z: tau * tf.reduce_sum(huber(z), -1)
    all_embed = tf.get_variable('all_embed', [1024, int(embed1.shape[-1])],
                                initializer=tf.zeros_initializer())
    upd = all_embed.assign(tf.concat([all_embed[batch_size:], embed2], 0))
    with tf.control_dependencies([upd]):
      close = 1 * tf.reduce_mean(fn(embed1 + action_embed - embed2))
      prior_log_probs = tf.reduce_logsumexp(
          -fn((embed1 + action_embed)[:, None, :] - all_embed[None, :, :]),
          axis=-1) - tf.log(tf.to_float(all_embed.shape[0]))
      far = tf.reduce_mean(tf.exp(-fn((embed1 + action_embed)[1:] - embed2[:-1])
                                  - tf.stop_gradient(prior_log_probs[1:])))
      repr_log_probs = tf.stop_gradient(
          -fn(embed1 + action_embed - embed2) - prior_log_probs) / tau
    return close + far, repr_log_probs, indices
//  [python]
def loss(self, states, next_states, low_actions, low_states):
    batch_size = tf.shape(states)[0]
    d = int(low_states.shape[1])
    # Sample indices into meta-transition to train on.
    probs = 0.99 ** tf.range(d, dtype=tf.float32)
    probs *= tf.constant([1.0] * (d - 1) + [1.0 / (1 - 0.99)],
                         dtype=tf.float32)
    probs /= tf.reduce_sum(probs)
    index_dist = tf.distributions.Categorical(probs=probs, dtype=tf.int64)
    indices = index_dist.sample(batch_size)
    batch_size = tf.cast(batch_size, tf.int64)
    next_indices = tf.concat(
        [tf.range(batch_size, dtype=tf.int64)[:, None],
         (1 + indices[:, None]) % d], -1)
    new_next_states = tf.where(indices < d - 1,
                               tf.gather_nd(low_states, next_indices),
                               next_states)
    next_states = new_next_states

    embed1 = tf.to_float(self._state_preprocess_net(states))
    embed2 = tf.to_float(self._state_preprocess_net(next_states))
    action_embed = self._action_embed_net(
        tf.layers.flatten(low_actions), states=states)

    tau = 2.0
    fn = lambda z: tau * tf.reduce_sum(huber(z), -1)
    all_embed = tf.get_variable('all_embed', [1024, int(embed1.shape[-1])],
                                initializer=tf.zeros_initializer())
    upd = all_embed.assign(tf.concat([all_embed[batch_size:], embed2], 0))
    with tf.control_dependencies([upd]):
      close = 1 * tf.reduce_mean(fn(embed1 + action_embed - embed2))
      prior_log_probs = tf.reduce_logsumexp(
          -fn((embed1 + action_embed)[:, None, :] - all_embed[None, :, :]),
          axis=-1) - tf.log(tf.to_float(all_embed.shape[0]))
      far = tf.reduce_mean(tf.exp(-fn((embed1 + action_embed)[1:] - embed2[:-1])
                                  - tf.stop_gradient(prior_log_probs[1:])))
      repr_log_probs = tf.stop_gradient(
          -fn(embed1 + action_embed - embed2) - prior_log_probs) / tau
    return close + far, repr_log_probs, indices
//  [python]
def create_padding_mask(decoder_token_ids):
    """
    Creates a matrix mask for the padding cells
    
    Arguments:
        decoder_token_ids -- (n, m) matrix
    
    Returns:
        mask -- (n, 1, 1, m) binary tensor
    """    
    seq = 1 - tf.to_float(tf.math.equal(decoder_token_ids, 0))
  
    # add extra dimensions to add the padding
    # to the attention logits.
    return seq[:, tf.newaxis, :]
//  [python]
def create_padding_mask(decoder_token_ids):
    """
    Creates a matrix mask for the padding cells
    
    Arguments:
        decoder_token_ids -- (n, m) matrix
    
    Returns:
        mask -- (n, 1, 1, m) binary tensor
    """    
    seq = 1 - tf.cast(tf.math.equal(decoder_token_ids, 0), tf.float32)
  
    # add extra dimensions to add the padding
    # to the attention logits.
    return seq[:, tf.newaxis, :]
//  [python]
def cond_begin_episode_op(self, cond, input_vars, mode, meta_action_fn):
    """Returns op that resets agent at beginning of episodes.

    A new episode is begun if the cond op evalues to `False`.

    Args:
      cond: a Boolean tensor variable.
      input_vars: A list of tensor variables.
      mode: a string representing the mode=[train, explore, eval].
    Returns:
      Conditional begin op.
    """
    (state, action, reward, next_state,
     state_repr, next_state_repr) = input_vars
    def continue_fn():
      """Continue op fn."""
      items = [state, action, reward, next_state,
               state_repr, next_state_repr] + list(self.context_vars)
      batch_items = [tf.expand_dims(item, 0) for item in items]
      (states, actions, rewards, next_states,
       state_reprs, next_state_reprs) = batch_items[:6]
      context_reward = self.compute_rewards(
          mode, state_reprs, actions, rewards, next_state_reprs,
          batch_items[6:])[0][0]
      context_reward = tf.cast(context_reward, dtype=reward.dtype)
      if self.meta_agent is not None:
        meta_action = tf.concat(self.context_vars, -1)
        items = [state, meta_action, reward, next_state,
                 state_repr, next_state_repr] + list(self.meta_agent.context_vars)
        batch_items = [tf.expand_dims(item, 0) for item in items]
        (states, meta_actions, rewards, next_states,
         state_reprs, next_state_reprs) = batch_items[:6]
        meta_reward = self.meta_agent.compute_rewards(
            mode, states, meta_actions, rewards,
            next_states, batch_items[6:])[0][0]
        meta_reward = tf.cast(meta_reward, dtype=reward.dtype)
      else:
        meta_reward = tf.constant(0, dtype=reward.dtype)

      with tf.control_dependencies([context_reward, meta_reward]):
        step_ops = self.tf_context.step(mode=mode, agent=self._meta_agent,
                                        state=state,
                                        next_state=next_state,
                                        state_repr=state_repr,
                                        next_state_repr=next_state_repr,
                                        action_fn=meta_action_fn)
      with tf.control_dependencies(step_ops):
        context_reward, meta_reward = map(tf.identity, [context_reward, meta_reward])
      return context_reward, meta_reward
    def begin_episode_fn():
      """Begin op fn."""
      begin_ops = self.begin_episode_ops(mode=mode, action_fn=meta_action_fn, state=state)
      with tf.control_dependencies(begin_ops):
        return tf.zeros_like(reward), tf.zeros_like(reward)
    with tf.control_dependencies(input_vars):
      cond_begin_episode_op = tf.cond(cond, continue_fn, begin_episode_fn)
    return cond_begin_episode_op
//  [python]
def positional_encoding(positions, d):
    """
    Precomputes a matrix with all the positional encodings 
    
    Arguments:
        positions (int) -- Maximum number of positions to be encoded 
        d (int) -- Encoding size 
    
    Returns:
        pos_encoding -- (1, position, d_model) A matrix with the positional encodings
    """
    # START CODE HERE
    # initialize a matrix angle_rads of all the angles
    angle_rads = get_angles(np.arange(positions)[:, np.newaxis],
                            np.arange(d)[np.newaxis, :],
                            d)
  
    # apply sin to even indices in the array; 2i
    angle_rads[:, 0::2] = np.sin(angle_rads[:, 0::2])
  
    # apply cos to odd indices in the array; 2i+1
    angle_rads[:, 1::2] = np.cos(angle_rads[:, 1::2])
    # END CODE HERE
    
    pos_encoding = angle_rads[np.newaxis, ...]
    
    return t.cast(pos_encoding, dtype=t.float32)
//  [python]
def positional_encoding(positions, d):
    """
    Precomputes a matrix with all the positional encodings 
    
    Arguments:
        positions (int) -- Maximum number of positions to be encoded 
        d (int) -- Encoding size 
    
    Returns:
        pos_encoding -- (1, position, d_model) A matrix with the positional encodings
    """
    # START CODE HERE
    # initialize a matrix angle_rads of all the angles
    angle_rads = get_angles(np.arange(positions)[:, np.newaxis],
                            np.arange(d)[np.newaxis, :],
                            d)
  
    # apply sin to even indices in the array; 2i
    angle_rads[:, 0::2] = np.sin(angle_rads[:, 0::2])
  
    # apply cos to odd indices in the array; 2i+1
    angle_rads[:, 1::2] = np.cos(angle_rads[:, 1::2])
    # END CODE HERE
    
    pos_encoding = angle_rads[np.newaxis, ...]
    
    return tf.cast(pos_encoding, dtype=tf.float32)
//  [python]
def call(self, x, training, mask):
       """
       Forward pass for the Encoder
       
       Arguments:
           x -- Tensor of shape (batch_size, input_seq_len)
           training -- Boolean, set to true to activate
                       the training mode for dropout layers
           mask -- Boolean mask to ensure that the padding is not 
                   treated as part of the input
       Returns:
           out2 -- Tensor of shape (batch_size, input_seq_len, fully_connected_dim)
       """
       #mask = create_padding_mask(x)
       seq_len = tf.shape(x)[1]
       
       # START CODE HERE
       # Pass input through the Embedding layer
       x = self.embedding(x)  # (batch_size, input_seq_len, fully_connected_dim)
       # Scale embedding by multiplying it by the square root of the embedding dimension
       x *= tf.sqrt(tf.to_float(self.embedding_dim))
       # Add the position encoding to embedding
       x += self.pos_encoding[:, :seq_len, :]
       # Pass the encoded embedding through a dropout layer
       x = self.dropout(x, training = training)
       # Pass the output through the stack of encoding layers 
       for i in range(self.num_layers):
           x = self.enc_layers[i](x, training, mask)
       # END CODE HERE

       return x
//  [python]
def call(self, x, training, mask):
       """
       Forward pass for the Encoder
       
       Arguments:
           x -- Tensor of shape (batch_size, input_seq_len)
           training -- Boolean, set to true to activate
                       the training mode for dropout layers
           mask -- Boolean mask to ensure that the padding is not 
                   treated as part of the input
       Returns:
           out2 -- Tensor of shape (batch_size, input_seq_len, fully_connected_dim)
       """
       #mask = create_padding_mask(x)
       seq_len = tf.shape(x)[1]
       
       # START CODE HERE
       # Pass input through the Embedding layer
       x = self.embedding(x)  # (batch_size, input_seq_len, fully_connected_dim)
       # Scale embedding by multiplying it by the square root of the embedding dimension
       x *= tf.math.sqrt(tf.cast(self.embedding_dim, tf.float32))
       # Add the position encoding to embedding
       x += self.pos_encoding[:, :seq_len, :]
       # Pass the encoded embedding through a dropout layer
       x = self.dropout(x, training = training)
       # Pass the output through the stack of encoding layers 
       for i in range(self.num_layers):
           x = self.enc_layers[i](x, training, mask)
       # END CODE HERE

       return x
//  [python]
def __init__(self, clip_percentile: int, history_size: int = 10000):
        self._clip_percentile = tf.cast(clip_percentile, tf.float64)
        self._grad_history = tf.Variable(tf.zeros(history_size), trainable=False)
        self._index = tf.Variable(0, trainable=False)
        self._history_size = history_size
//  [python]
def compute_accuracy(v_xs,v_ys,sess):
    #prediction 变为全剧变量
    global prediction
    y_pre = sess.run(prediction,feed_dict={xs:v_xs})
    #预测值每行是10列，tf.argmax(数据，axis），相等为1，不想等为0
    correct_prediction = tf.equal(tf.argmax(y_pre,1),tf.argmax(v_ys,1))
    # 计算平均值，即计算准确率
    accuracy = tf.reduce_mean(tf.to_float(correct_prediction))
    # 运行我们的accuracy这一步
    result = sess.run(accuracy,feed_dict={xs:v_xs,ys:v_ys})
    return result
//  [python]
def compute_accuracy(v_xs,v_ys,sess):
    #prediction 变为全剧变量
    global prediction
    y_pre = sess.run(prediction,feed_dict={xs:v_xs})
    #预测值每行是10列，tf.argmax(数据，axis），相等为1，不想等为0
    correct_prediction = tf.equal(tf.argmax(y_pre,1),tf.argmax(v_ys,1))
    # 计算平均值，即计算准确率
    accuracy = tf.reduce_mean(tf.cast(correct_prediction,tf.float32))
    # 运行我们的accuracy这一步
    result = sess.run(accuracy,feed_dict={xs:v_xs,ys:v_ys})
    return result
//  [python]
def scaled_dot_product_attention(q, k, v, mask):
    """
    Calculate the attention weights.
      q, k, v must have matching leading dimensions.
      k, v must have matching penultimate dimension, i.e.: seq_len_k = seq_len_v.
      The mask has different shapes depending on its type(padding or look ahead) 
      but it must be broadcastable for addition.

    Arguments:
        q -- query shape == (..., seq_len_q, depth)
        k -- key shape == (..., seq_len_k, depth)
        v -- value shape == (..., seq_len_v, depth_v)
        mask: Float tensor with shape broadcastable 
              to (..., seq_len_q, seq_len_k). Defaults to None.

    Returns:
        output -- attention_weights
    """
   # START CODE HERE
    
    matmul_qk = tf.matmul(q, k, transpose_b = True)  # (..., seq_len_q, seq_len_k)

    # scale matmul_qk
    dk = tf.to_float(tf.shape(k)[-1])
    scaled_attention_logits = matmul_qk / tf.math.sqrt(dk)

    # add the mask to the scaled tensor.
    if mask is not None: # Don't replace this None
        scaled_attention_logits += ((1 - mask) * -1e9)

    # softmax is normalized on the last axis (seq_len_k) so that the scores
    # add up to 1.
    attention_weights = tf.nn.softmax(scaled_attention_logits, axis = -1)  # (..., seq_len_q, seq_len_k)

    output = tf.matmul(attention_weights, v)  # (..., seq_len_q, depth_v)
    
    # END CODE HERE
    return output, attention_weights
//  [python]
def scaled_dot_product_attention(q, k, v, mask):
    """
    Calculate the attention weights.
      q, k, v must have matching leading dimensions.
      k, v must have matching penultimate dimension, i.e.: seq_len_k = seq_len_v.
      The mask has different shapes depending on its type(padding or look ahead) 
      but it must be broadcastable for addition.

    Arguments:
        q -- query shape == (..., seq_len_q, depth)
        k -- key shape == (..., seq_len_k, depth)
        v -- value shape == (..., seq_len_v, depth_v)
        mask: Float tensor with shape broadcastable 
              to (..., seq_len_q, seq_len_k). Defaults to None.

    Returns:
        output -- attention_weights
    """
   # START CODE HERE
    
    matmul_qk = tf.matmul(q, k, transpose_b = True)  # (..., seq_len_q, seq_len_k)

    # scale matmul_qk
    dk = tf.cast(tf.shape(k)[-1], tf.float32)
    scaled_attention_logits = matmul_qk / tf.math.sqrt(dk)

    # add the mask to the scaled tensor.
    if mask is not None: # Don't replace this None
        scaled_attention_logits += ((1 - mask) * -1e9)

    # softmax is normalized on the last axis (seq_len_k) so that the scores
    # add up to 1.
    attention_weights = tf.nn.softmax(scaled_attention_logits, axis = -1)  # (..., seq_len_q, seq_len_k)

    output = tf.matmul(attention_weights, v)  # (..., seq_len_q, depth_v)
    
    # END CODE HERE
    return output, attention_weights
//  [python]
def __call__(self, logits, samples, logit_length=None):
        target = samples["output"]
        shape = tf.shape(logits)
        target = tf.reshape(target, shape)
        loss = target - logits
        # mpc mask
        mask = tf.to_float(tf.math.equal(tf.reshape(samples["input"], shape), 0))
        loss *= mask
        # sequence length mask
        seq_mask = tf.sequence_mask(logit_length, shape[1], dtype=loss.dtype)
        seq_mask = tf.tile(seq_mask[:, :, tf.newaxis], [1, 1, shape[2]])
        loss *= seq_mask
        loss = tf.reduce_sum(tf.abs(loss, name="L1_loss"), axis=-1)
        loss = tf.reduce_mean(loss)
        return loss
//  [python]
def __call__(self, logits, samples, logit_length=None):
        target = samples["output"]
        shape = tf.shape(logits)
        target = tf.reshape(target, shape)
        loss = target - logits
        # mpc mask
        mask = tf.cast(tf.math.equal(tf.reshape(samples["input"], shape), 0), loss.dtype)
        loss *= mask
        # sequence length mask
        seq_mask = tf.sequence_mask(logit_length, shape[1], dtype=loss.dtype)
        seq_mask = tf.tile(seq_mask[:, :, tf.newaxis], [1, 1, shape[2]])
        loss *= seq_mask
        loss = tf.reduce_sum(tf.abs(loss, name="L1_loss"), axis=-1)
        loss = tf.reduce_mean(loss)
        return loss
//  [python]
def __call__(self, logits, samples, logit_length=None):
        labels = insert_eos_in_labels(samples["output"], self.eos, samples["output_length"])
        mask = tf.math.logical_not(tf.math.equal(labels, 0))
        labels = tf.one_hot(indices=labels, depth=self.num_classes)
        if self.smoothing_rate != 0.0:
            labels = apply_label_smoothing(labels, 
                num_classes=self.num_classes, 
                smoothing_rate=self.smoothing_rate)
        seq_len = tf.shape(labels)[1]
        logits = logits[:, :seq_len, :]
        loss = self.call(labels, logits)
        mask = tf.to_float(mask)
        loss *= mask
        if self.by_token:
            return tf.divide(tf.reduce_sum(loss), tf.reduce_sum(mask))
        if self.by_sequence:
            loss = tf.reduce_sum(loss, axis=-1)
        return tf.reduce_mean(loss)
//  [python]
def __call__(self, logits, samples, logit_length=None):
        labels = insert_eos_in_labels(samples["output"], self.eos, samples["output_length"])
        mask = tf.math.logical_not(tf.math.equal(labels, 0))
        labels = tf.one_hot(indices=labels, depth=self.num_classes)
        if self.smoothing_rate != 0.0:
            labels = apply_label_smoothing(labels, 
                num_classes=self.num_classes, 
                smoothing_rate=self.smoothing_rate)
        seq_len = tf.shape(labels)[1]
        logits = logits[:, :seq_len, :]
        loss = self.call(labels, logits)
        mask = tf.cast(mask, dtype=loss.dtype)
        loss *= mask
        if self.by_token:
            return tf.divide(tf.reduce_sum(loss), tf.reduce_sum(mask))
        if self.by_sequence:
            loss = tf.reduce_sum(loss, axis=-1)
        return tf.reduce_mean(loss)
//  [python]
def call(self, x, enc_output, training, 
           look_ahead_mask, padding_mask):
        """
        Forward  pass for the Decoder
        
        Arguments:
            x -- Tensor of shape (batch_size, target_seq_len, fully_connected_dim)
            enc_output --  Tensor of shape(batch_size, input_seq_len, fully_connected_dim)
            training -- Boolean, set to true to activate
                        the training mode for dropout layers
            look_ahead_mask -- Boolean mask for the target_input
            padding_mask -- Boolean mask for the second multihead attention layer
        Returns:
            x -- Tensor of shape (batch_size, target_seq_len, fully_connected_dim)
            attention_weights - Dictionary of tensors containing all the attention weights
                                each of shape Tensor of shape (batch_size, num_heads, target_seq_len, input_seq_len)
        """

        seq_len = tf.shape(x)[1]
        attention_weights = {}
        
         # START CODE HERE
        # create word embeddings 
        x = self.embedding(x)  # (batch_size, target_seq_len, fully_connected_dim)
        
        # scale embeddings by multiplying by the square root of their dimension
        x *= tf.math.sqrt(tf.to_float(self.embedding_dim))
        
        # calculate positional encodings and add to word embedding
        x += self.pos_encoding[:, :seq_len, :]

        # apply a dropout layer to x
        x = self.dropout(x, training = training)

        # use a for loop to pass x through a stack of decoder layers and update attention_weights (~4 lines total)
        for i in range(self.num_layers):
            # pass x and the encoder output through a stack of decoder layers and save the attention weights
            # of block 1 and 2 (~1 line)
            x, block1, block2 = self.dec_layers[i](x, enc_output, training,
                                                 look_ahead_mask, padding_mask)

            #update attention_weights dictionary with the attention weights of block 1 and block 2
            attention_weights['decoder_layer{}_block1_self_att'.format(i+1)] = block1
            attention_weights['decoder_layer{}_block2_decenc_att'.format(i+1)] = block2
        # END CODE HERE
        
        
        # x.shape == (batch_size, target_seq_len, fully_connected_dim)
        return x, attention_weights