jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def __init__(self, config):

        self.config = config
        self.lr = config["lr"]
        self.char_dim = config["char_dim"]
        self.lstm_dim = config["lstm_dim"]
        self.seg_dim = config["seg_dim"]

        self.num_tags = config["num_tags"]
        self.num_chars = config["num_chars"]
        self.num_segs = 4

        self.global_step = tf.Variable(0, trainable=False)
        self.best_dev_f1 = tf.Variable(0.0, trainable=False)
        self.best_test_f1 = tf.Variable(0.0, trainable=False)
        self.initializer = initializers.xavier_initializer()

        # add placeholders for the model

        self.char_inputs = tf.placeholder(dtype=tf.int32,
                                          shape=[None, None],
                                          name="ChatInputs")
        self.seg_inputs = tf.placeholder(dtype=tf.int32,
                                         shape=[None, None],
                                         name="SegInputs")

        self.targets = tf.placeholder(dtype=tf.int32,
                                      shape=[None, None],
                                      name="Targets")
        # dropout keep prob
        self.dropout = tf.placeholder(dtype=tf.float32,
                                      name="Dropout")

        used = tf.sign(tf.abs(self.char_inputs))
        length = tf.reduce_sum(used, reduction_indices=1)
        self.lengths = tf.to_int32(length)
        self.batch_size = tf.shape(self.char_inputs)[0]
        self.num_steps = tf.shape(self.char_inputs)[-1]

        # embeddings for chinese character and segmentation representation
        embedding = self.embedding_layer(self.char_inputs, self.seg_inputs, config)

        # apply dropout before feed to lstm layer
        lstm_inputs = tf.nn.dropout(embedding, self.dropout)

        # bi-directional lstm layer
        lstm_outputs = self.biLSTM_layer(lstm_inputs, self.lstm_dim, self.lengths)

        # logits for tags
        self.logits = self.project_layer(lstm_outputs)

        # loss of the model
        self.loss = self.loss_layer(self.logits, self.lengths)

        with tf.variable_scope("optimizer"):
            optimizer = self.config["optimizer"]
            if optimizer == "sgd":
                self.opt = tf.train.GradientDescentOptimizer(self.lr)
            elif optimizer == "adam":
                self.opt = tf.train.AdamOptimizer(self.lr)
            elif optimizer == "adgrad":
                self.opt = tf.train.AdagradOptimizer(self.lr)
            else:
                raise KeyError

            # apply grad clip to avoid gradient explosion
            grads_vars = self.opt.compute_gradients(self.loss)
            capped_grads_vars = [[tf.clip_by_value(g, -self.config["clip"], self.config["clip"]), v]
                                 for g, v in grads_vars]
            self.train_op = self.opt.apply_gradients(capped_grads_vars, self.global_step)

        # saver of the model
        self.saver = tf.train.Saver(tf.global_variables(), max_to_keep=5)
//  [python]
def __init__(self, config):

        self.config = config
        self.lr = config["lr"]
        self.char_dim = config["char_dim"]
        self.lstm_dim = config["lstm_dim"]
        self.seg_dim = config["seg_dim"]

        self.num_tags = config["num_tags"]
        self.num_chars = config["num_chars"]
        self.num_segs = 4

        self.global_step = tf.Variable(0, trainable=False)
        self.best_dev_f1 = tf.Variable(0.0, trainable=False)
        self.best_test_f1 = tf.Variable(0.0, trainable=False)
        self.initializer = initializers.xavier_initializer()

        # add placeholders for the model

        self.char_inputs = tf.placeholder(dtype=tf.int32,
                                          shape=[None, None],
                                          name="ChatInputs")
        self.seg_inputs = tf.placeholder(dtype=tf.int32,
                                         shape=[None, None],
                                         name="SegInputs")

        self.targets = tf.placeholder(dtype=tf.int32,
                                      shape=[None, None],
                                      name="Targets")
        # dropout keep prob
        self.dropout = tf.placeholder(dtype=tf.float32,
                                      name="Dropout")

        used = tf.sign(tf.abs(self.char_inputs))
        length = tf.reduce_sum(used, reduction_indices=1)
        self.lengths = tf.cast(length, tf.int32)
        self.batch_size = tf.shape(self.char_inputs)[0]
        self.num_steps = tf.shape(self.char_inputs)[-1]

        # embeddings for chinese character and segmentation representation
        embedding = self.embedding_layer(self.char_inputs, self.seg_inputs, config)

        # apply dropout before feed to lstm layer
        lstm_inputs = tf.nn.dropout(embedding, self.dropout)

        # bi-directional lstm layer
        lstm_outputs = self.biLSTM_layer(lstm_inputs, self.lstm_dim, self.lengths)

        # logits for tags
        self.logits = self.project_layer(lstm_outputs)

        # loss of the model
        self.loss = self.loss_layer(self.logits, self.lengths)

        with tf.variable_scope("optimizer"):
            optimizer = self.config["optimizer"]
            if optimizer == "sgd":
                self.opt = tf.train.GradientDescentOptimizer(self.lr)
            elif optimizer == "adam":
                self.opt = tf.train.AdamOptimizer(self.lr)
            elif optimizer == "adgrad":
                self.opt = tf.train.AdagradOptimizer(self.lr)
            else:
                raise KeyError

            # apply grad clip to avoid gradient explosion
            grads_vars = self.opt.compute_gradients(self.loss)
            capped_grads_vars = [[tf.clip_by_value(g, -self.config["clip"], self.config["clip"]), v]
                                 for g, v in grads_vars]
            self.train_op = self.opt.apply_gradients(capped_grads_vars, self.global_step)

        # saver of the model
        self.saver = tf.train.Saver(tf.global_variables(), max_to_keep=5)
//  [python]
def make_anchors(self, feature_pyramid):
        with tf.variable_scope('make_anchors'):
            anchor_list = []
            level_list = cfgs.LEVEL
            with tf.name_scope('make_anchors_all_level'):
                for level, base_anchor_size, stride in zip(level_list, cfgs.BASE_ANCHOR_SIZE_LIST, cfgs.ANCHOR_STRIDE):
                    '''
                    (level, base_anchor_size) tuple:
                    (P3, 32), (P4, 64), (P5, 128), (P6, 256), (P7, 512)
                    '''
                    featuremap_height, featuremap_width = tf.shape(feature_pyramid[level])[1], \
                                                          tf.shape(feature_pyramid[level])[2]

                    featuremap_height = tf.to_float(featuremap_height)
                    featuremap_width = tf.to_float(featuremap_width)

                    if self.method == 'H':
                        tmp_anchors = tf.py_func(generate_anchors.generate_anchors_pre,
                                                 inp=[featuremap_height, featuremap_width, stride,
                                                      np.array(cfgs.ANCHOR_SCALES) * stride, cfgs.ANCHOR_RATIOS, 4.0],
                                                 Tout=[tf.float32])

                        tmp_anchors = tf.reshape(tmp_anchors, [-1, 4])
                    else:
                        tmp_anchors = generate_rotate_anchors.make_anchors(base_anchor_size=base_anchor_size,
                                                                           anchor_scales=cfgs.ANCHOR_SCALES,
                                                                           anchor_ratios=cfgs.ANCHOR_RATIOS,
                                                                           anchor_angles=cfgs.ANCHOR_ANGLES,
                                                                           featuremap_height=featuremap_height,
                                                                           featuremap_width=featuremap_width,
                                                                           stride=stride)
                        tmp_anchors = tf.reshape(tmp_anchors, [-1, 5])

                    anchor_list.append(tmp_anchors)

                # all_level_anchors = tf.concat(anchor_list, axis=0)
            return anchor_list
//  [python]
def make_anchors(self, feature_pyramid):
        with tf.variable_scope('make_anchors'):
            anchor_list = []
            level_list = cfgs.LEVEL
            with tf.name_scope('make_anchors_all_level'):
                for level, base_anchor_size, stride in zip(level_list, cfgs.BASE_ANCHOR_SIZE_LIST, cfgs.ANCHOR_STRIDE):
                    '''
                    (level, base_anchor_size) tuple:
                    (P3, 32), (P4, 64), (P5, 128), (P6, 256), (P7, 512)
                    '''
                    featuremap_height, featuremap_width = tf.shape(feature_pyramid[level])[1], \
                                                          tf.shape(feature_pyramid[level])[2]

                    featuremap_height = tf.cast(featuremap_height, tf.float32)
                    featuremap_width = tf.cast(featuremap_width, tf.float32)

                    if self.method == 'H':
                        tmp_anchors = tf.py_func(generate_anchors.generate_anchors_pre,
                                                 inp=[featuremap_height, featuremap_width, stride,
                                                      np.array(cfgs.ANCHOR_SCALES) * stride, cfgs.ANCHOR_RATIOS, 4.0],
                                                 Tout=[tf.float32])

                        tmp_anchors = tf.reshape(tmp_anchors, [-1, 4])
                    else:
                        tmp_anchors = generate_rotate_anchors.make_anchors(base_anchor_size=base_anchor_size,
                                                                           anchor_scales=cfgs.ANCHOR_SCALES,
                                                                           anchor_ratios=cfgs.ANCHOR_RATIOS,
                                                                           anchor_angles=cfgs.ANCHOR_ANGLES,
                                                                           featuremap_height=featuremap_height,
                                                                           featuremap_width=featuremap_width,
                                                                           stride=stride)
                        tmp_anchors = tf.reshape(tmp_anchors, [-1, 5])

                    anchor_list.append(tmp_anchors)

                # all_level_anchors = tf.concat(anchor_list, axis=0)
            return anchor_list
//  [python]
def call(
        self,
        x: tf.Tensor,
        target: Optional[List[str]] = None,
        return_model_output: bool = False,
        return_preds: bool = False,
        **kwargs: Any,
    ) -> Dict[str, Any]:
        features = self.feat_extractor(x, **kwargs)
        # vertical max pooling --> (N, C, W)
        pooled_features = tf.reduce_max(features, axis=1)
        # holistic (N, C)
        encoded = self.encoder(pooled_features, **kwargs)

        if target is not None:
            gt, seq_len = self.build_target(target)
            seq_len = t.cast(seq_len, tf.int32)

        if kwargs.get("training", False) and target is None:
            raise ValueError("Need to provide labels during training for teacher forcing")

        decoded_features = _bf16_to_float32(
            self.decoder(features, encoded, gt=None if target is None else gt, **kwargs)
        )

        out: Dict[str, tf.Tensor] = {}
        if self.exportable:
            out["logits"] = decoded_features
            return out

        if return_model_output:
            out["out_map"] = decoded_features

        if target is None or return_preds:
            # Post-process boxes
            out["preds"] = self.postprocessor(decoded_features)

        if target is not None:
            out["loss"] = self.compute_loss(decoded_features, gt, seq_len)

        return out
//  [python]
def call(
        self,
        x: tf.Tensor,
        target: Optional[List[str]] = None,
        return_model_output: bool = False,
        return_preds: bool = False,
        **kwargs: Any,
    ) -> Dict[str, Any]:
        features = self.feat_extractor(x, **kwargs)
        # vertical max pooling --> (N, C, W)
        pooled_features = tf.reduce_max(features, axis=1)
        # holistic (N, C)
        encoded = self.encoder(pooled_features, **kwargs)

        if target is not None:
            gt, seq_len = self.build_target(target)
            seq_len = tf.cast(seq_len, tf.int32)

        if kwargs.get("training", False) and target is None:
            raise ValueError("Need to provide labels during training for teacher forcing")

        decoded_features = _bf16_to_float32(
            self.decoder(features, encoded, gt=None if target is None else gt, **kwargs)
        )

        out: Dict[str, tf.Tensor] = {}
        if self.exportable:
            out["logits"] = decoded_features
            return out

        if return_model_output:
            out["out_map"] = decoded_features

        if target is None or return_preds:
            # Post-process boxes
            out["preds"] = self.postprocessor(decoded_features)

        if target is not None:
            out["loss"] = self.compute_loss(decoded_features, gt, seq_len)

        return out
//  [python]
def _compute_top_k_loss(self, loss):
    """Computs top k loss."""
    batch_size = tf.shape(loss)[0]
    loss = tf.reshape(loss, shape=[batch_size, -1])

    top_k_pixels = tf.to_int32(
        self._top_k_percent_pixels *
        tf.to_float(tf.shape(loss)[-1]),
        dtype=tf.int32)

    # shape: [batch_size, top_k_pixels]
    per_sample_top_k_loss = tf.map_fn(
        fn=lambda x: tf.nn.top_k(x, k=top_k_pixels, sorted=False)[0],
        elems=loss,
        parallel_iterations=32,
        fn_output_signature=tf.float32)

    # shape: [batch_size]
    per_sample_normalizer = tf.reduce_sum(
        tf.to_float(
            tf.not_equal(per_sample_top_k_loss, 0.0)),
        axis=-1) + EPSILON
    per_sample_normalized_loss = tf.reduce_sum(
        per_sample_top_k_loss, axis=-1) / per_sample_normalizer

    normalized_loss = tf_utils.safe_mean(per_sample_normalized_loss)
    return normalized_loss
//  [python]
def _compute_top_k_loss(self, loss):
    """Computs top k loss."""
    batch_size = tf.shape(loss)[0]
    loss = tf.reshape(loss, shape=[batch_size, -1])

    top_k_pixels = tf.cast(
        self._top_k_percent_pixels *
        tf.cast(tf.shape(loss)[-1], dtype=tf.float32),
        dtype=tf.int32)

    # shape: [batch_size, top_k_pixels]
    per_sample_top_k_loss = tf.map_fn(
        fn=lambda x: tf.nn.top_k(x, k=top_k_pixels, sorted=False)[0],
        elems=loss,
        parallel_iterations=32,
        fn_output_signature=tf.float32)

    # shape: [batch_size]
    per_sample_normalizer = tf.reduce_sum(
        tf.cast(
            tf.not_equal(per_sample_top_k_loss, 0.0),
            dtype=tf.float32),
        axis=-1) + EPSILON
    per_sample_normalized_loss = tf.reduce_sum(
        per_sample_top_k_loss, axis=-1) / per_sample_normalizer

    normalized_loss = tf_utils.safe_mean(per_sample_normalized_loss)
    return normalized_loss
//  [python]
def get_gaussian_pdf_equal_split_points(k):
  tf.assert_greater(k, 2)
  lp = t.cast(tf.linspace(0, 1, k+1), dtype=tf.float32)
  pk = tfd.Normal(0.0, 1.0).quantile(lp[1:-1])
  qk_pre = (pk[:-1] + pk[1:])/2
  qk_first = 2*pk[0] - qk_pre[0]
  qk_last = 2*pk[-1] - qk_pre[-1]
  qk = tf.concat([[qk_first], qk_pre, [qk_last]], axis=0)
  return qk
//  [python]
def get_gaussian_pdf_equal_split_points(k):
  tf.assert_greater(k, 2)
  lp = tf.cast(tf.linspace(0, 1, k+1), dtype=tf.float32)
  pk = tfd.Normal(0.0, 1.0).quantile(lp[1:-1])
  qk_pre = (pk[:-1] + pk[1:])/2
  qk_first = 2*pk[0] - qk_pre[0]
  qk_last = 2*pk[-1] - qk_pre[-1]
  qk = tf.concat([[qk_first], qk_pre, [qk_last]], axis=0)
  return qk
//  [python]
def top_grasp_acc_summaries(ops, thres=[0.62,0.65,0.68]):
    """
    Calculates the accuracy of grasp contact predictions at different thresholds

    Arguments:
        ops {dict} -- dict of tf tensors

    Keyword Arguments:
        thres {list} -- thresholds to report (default: {[0.62,0.65,0.68]})

    Returns:
        [tf.variable] -- summary update ops
    """
    
    binary_seg = tf.reshape(ops['binary_seg_pred'],[ops['binary_seg_pred'].shape[0],ops['binary_seg_pred'].shape[1]])
    max_conf_idcs = tf.argmax(binary_seg, axis=1)
    max_confs = tf.math.reduce_max(binary_seg, axis=1)

    row_idcs = tf.to_int64(tf.range(ops['grasp_suc_labels_pc'].shape[0]))
    nd_gather_idcs = tf.stack((row_idcs, max_conf_idcs), axis=1)
    top_labels = tf.gather_nd(ops['grasp_suc_labels_pc'], nd_gather_idcs)

    acc_update_ops = [0] * len(thres)
    accs = [0] * len(thres)
    for i,thr in enumerate(thres):
        accs[i], acc_update_ops[i] = tf.metrics.accuracy(top_labels, tf.greater(max_confs, tf.constant(thr)))
        tf.summary.scalar('Top Grasp Accuracy Thres %s' % thr, accs[i])

    return acc_update_ops
//  [python]
def top_grasp_acc_summaries(ops, thres=[0.62,0.65,0.68]):
    """
    Calculates the accuracy of grasp contact predictions at different thresholds

    Arguments:
        ops {dict} -- dict of tf tensors

    Keyword Arguments:
        thres {list} -- thresholds to report (default: {[0.62,0.65,0.68]})

    Returns:
        [tf.variable] -- summary update ops
    """
    
    binary_seg = tf.reshape(ops['binary_seg_pred'],[ops['binary_seg_pred'].shape[0],ops['binary_seg_pred'].shape[1]])
    max_conf_idcs = tf.argmax(binary_seg, axis=1)
    max_confs = tf.math.reduce_max(binary_seg, axis=1)

    row_idcs = tf.cast(tf.range(ops['grasp_suc_labels_pc'].shape[0]), tf.int64)
    nd_gather_idcs = tf.stack((row_idcs, max_conf_idcs), axis=1)
    top_labels = tf.gather_nd(ops['grasp_suc_labels_pc'], nd_gather_idcs)

    acc_update_ops = [0] * len(thres)
    accs = [0] * len(thres)
    for i,thr in enumerate(thres):
        accs[i], acc_update_ops[i] = tf.metrics.accuracy(top_labels, tf.greater(max_confs, tf.constant(thr)))
        tf.summary.scalar('Top Grasp Accuracy Thres %s' % thr, accs[i])

    return acc_update_ops
//  [python]
def generate_weighted_samples(heatmap, limits):
  """Transforms a 2D heatmap into 2D samples weighted by probabilities."""
  res = tf.shape(heatmap)
  unnorm_weights = tf.reshape(heatmap, (-1,))

  mesh = tf.meshgrid(tf.linspace(0, 1, res[1]), tf.linspace(0, 1, res[0]))
  samples_grid = tf.reshape(tf.stack(mesh, axis=-1), (-1, 2))[:, ::-1]

  scale = tf.convert_to_tensor([limits[0][1] - limits[0][0],
                                limits[1][1] - limits[1][0]])
  shift = tf.convert_to_tensor([limits[0][0], limits[1][0]])
  samples_grid, res, scale, shift = [
      tf.to_float(x) for x in [samples_grid, res, scale, shift]]
  samples = samples_grid * scale + shift
  weights = unnorm_weights / tf.reduce_sum(unnorm_weights)
  samples, weights = [tf.to_float(x) for x in [samples, weights]]
  return samples, weights
//  [python]
def generate_weighted_samples(heatmap, limits):
  """Transforms a 2D heatmap into 2D samples weighted by probabilities."""
  res = tf.shape(heatmap)
  unnorm_weights = tf.reshape(heatmap, (-1,))

  mesh = tf.meshgrid(tf.linspace(0, 1, res[1]), tf.linspace(0, 1, res[0]))
  samples_grid = tf.reshape(tf.stack(mesh, axis=-1), (-1, 2))[:, ::-1]

  scale = tf.convert_to_tensor([limits[0][1] - limits[0][0],
                                limits[1][1] - limits[1][0]])
  shift = tf.convert_to_tensor([limits[0][0], limits[1][0]])
  samples_grid, res, scale, shift = [
      tf.cast(x, tf.float32) for x in [samples_grid, res, scale, shift]]
  samples = samples_grid * scale + shift
  weights = unnorm_weights / tf.reduce_sum(unnorm_weights)
  samples, weights = [tf.cast(x, tf.float32) for x in [samples, weights]]
  return samples, weights
//  [python]
def compensate_heatmap_sigma(sigma, res, limits):
  scale = tf.convert_to_tensor([limits[0][1] - limits[0][0],
                                limits[1][1] - limits[1][0]])
  res = t.cast(res, tf.float32)
  pixel_noise_multiplier = 0.5
  noise_level = tf.reduce_max(scale / res) * pixel_noise_multiplier
  return tf.math.sqrt(sigma**2 + noise_level**2)
//  [python]
def compensate_heatmap_sigma(sigma, res, limits):
  scale = tf.convert_to_tensor([limits[0][1] - limits[0][0],
                                limits[1][1] - limits[1][0]])
  res = tf.cast(res, tf.float32)
  pixel_noise_multiplier = 0.5
  noise_level = tf.reduce_max(scale / res) * pixel_noise_multiplier
  return tf.math.sqrt(sigma**2 + noise_level**2)
//  [python]
def _parse_data(self, data, is_training):
    image = data['image']

    if self._augmenter is not None and is_training:
      image = self._augmenter.distort(image)

    image = preprocess_ops.normalize_image(image)

    category_mask = tf.cast(
        data['groundtruth_panoptic_category_mask'][:, :, 0],
        dtype=tf.float32)
    instance_mask = tf.cast(
        data['groundtruth_panoptic_instance_mask'][:, :, 0],
        dtype=tf.float32)

    # Flips image randomly during training.
    if self._aug_rand_hflip and is_training:
      masks = tf.stack([category_mask, instance_mask], axis=0)
      image, _, masks = preprocess_ops.random_horizontal_flip(
          image=image, masks=masks)
      category_mask = masks[0]
      instance_mask = masks[1]

    # Resizes and crops image.
    image, image_info = preprocess_ops.resize_and_crop_image(
        image,
        self._output_size,
        self._output_size,
        aug_scale_min=self._aug_scale_min if is_training else 1.0,
        aug_scale_max=self._aug_scale_max if is_training else 1.0)

    category_mask = self._resize_and_crop_mask(
        category_mask,
        image_info,
        is_training=is_training)
    instance_mask = self._resize_and_crop_mask(
        instance_mask,
        image_info,
        is_training=is_training)

    (instance_centers_heatmap,
     instance_centers_offset,
     semantic_weights) = self._encode_centers_and_offets(
         instance_mask=instance_mask[:, :, 0])

    # Cast image and labels as self._dtype
    image = tf.cast(image, dtype=self._dtype)
    category_mask = tf.cast(category_mask, dtype=self._dtype)
    instance_mask = tf.cast(instance_mask, dtype=self._dtype)
    instance_centers_heatmap = tf.cast(
        instance_centers_heatmap, dtype=self._dtype)
    instance_centers_offset = tf.cast(
        instance_centers_offset, dtype=self._dtype)

    valid_mask = tf.not_equal(
        category_mask, self._ignore_label)
    things_mask = tf.not_equal(
        instance_mask, self._ignore_label)

    labels = {
        'category_mask': category_mask,
        'instance_mask': instance_mask,
        'instance_centers_heatmap': instance_centers_heatmap,
        'instance_centers_offset': instance_centers_offset,
        'semantic_weights': semantic_weights,
        'valid_mask': valid_mask,
        'things_mask': things_mask,
        'image_info': image_info
    }
    return image, labels
//  [python]
def fast_fit_gmm(heatmap, limits, k_comp, steps=5, init_stddev=2.0):
  """Performs weighted EM clustering over the heatmap."""
  samples, weights = generate_weighted_samples(heatmap, limits)
  mus = init_means_with_pca(samples, weights, k_comp, init_stddev)
  sigmas = t.ones(k_comp)
  pis = t.ones(k_comp, dtype=t.float32) / t.cast(k_comp, t.float32)

  samples, weights, mus, sigmas, pis = [
      t.cast(x, t.float32) for x in [samples, weights, mus, sigmas, pis]]

  for _ in range(steps):
    assign_probs = compute_assignment(samples, mus, sigmas, pis)
    pis, mus, sigmas = estiamte_params(samples, weights, assign_probs)
    sigmas = compensate_heatmap_sigma(sigmas, heatmap.shape, limits)

  return pis, mus, sigmas
//  [python]
def fast_fit_gmm(heatmap, limits, k_comp, steps=5, init_stddev=2.0):
  """Performs weighted EM clustering over the heatmap."""
  samples, weights = generate_weighted_samples(heatmap, limits)
  mus = init_means_with_pca(samples, weights, k_comp, init_stddev)
  sigmas = tf.ones(k_comp)
  pis = tf.ones(k_comp, dtype=tf.float32) / tf.cast(k_comp, tf.float32)

  samples, weights, mus, sigmas, pis = [
      tf.cast(x, tf.float32) for x in [samples, weights, mus, sigmas, pis]]

  for _ in range(steps):
    assign_probs = compute_assignment(samples, mus, sigmas, pis)
    pis, mus, sigmas = estiamte_params(samples, weights, assign_probs)
    sigmas = compensate_heatmap_sigma(sigmas, heatmap.shape, limits)

  return pis, mus, sigmas
//  [python]
def estiamte_params(samples, sample_weights, assign_probs):
  """Estimate updated mixture parameters from weighted samples (M-step)."""
  total_sample_comp_weights = sample_weights[:, None] * assign_probs  # [N, k]

  new_nn_pis = sample_weights[:, None] * assign_probs  # [N, k]
  new_pis = tf.reduce_sum(new_nn_pis, axis=0) / tf.reduce_sum(new_nn_pis)

  # [N, k, d]
  weighted_samples = samples[:, None, :] * total_sample_comp_weights[:, :, None]
  new_means = (
      tf.reduce_sum(weighted_samples, axis=0)
      / (tf.reduce_sum(total_sample_comp_weights, axis=0)[:, None] + 1e-15))
  new_diffs = samples[:, None, :] - new_means[None, :, :]  # [N, k, d]
  new_dist_sq = tf.reduce_mean(tf.math.square(new_diffs), axis=2)  # [N, k]

  new_sigmas_sq = (
      tf.reduce_sum(new_dist_sq * total_sample_comp_weights, axis=0)
      / (tf.reduce_sum(total_sample_comp_weights, axis=0) + 1e-15)
  )
  new_sigmas = tf.math.sqrt(new_sigmas_sq)
  new_pis, new_means, new_sigmas = [
      tf.to_float(x) for x in [new_pis, new_means, new_sigmas]]
  return new_pis, new_means, new_sigmas
//  [python]
def estiamte_params(samples, sample_weights, assign_probs):
  """Estimate updated mixture parameters from weighted samples (M-step)."""
  total_sample_comp_weights = sample_weights[:, None] * assign_probs  # [N, k]

  new_nn_pis = sample_weights[:, None] * assign_probs  # [N, k]
  new_pis = tf.reduce_sum(new_nn_pis, axis=0) / tf.reduce_sum(new_nn_pis)

  # [N, k, d]
  weighted_samples = samples[:, None, :] * total_sample_comp_weights[:, :, None]
  new_means = (
      tf.reduce_sum(weighted_samples, axis=0)
      / (tf.reduce_sum(total_sample_comp_weights, axis=0)[:, None] + 1e-15))
  new_diffs = samples[:, None, :] - new_means[None, :, :]  # [N, k, d]
  new_dist_sq = tf.reduce_mean(tf.math.square(new_diffs), axis=2)  # [N, k]

  new_sigmas_sq = (
      tf.reduce_sum(new_dist_sq * total_sample_comp_weights, axis=0)
      / (tf.reduce_sum(total_sample_comp_weights, axis=0) + 1e-15)
  )
  new_sigmas = tf.math.sqrt(new_sigmas_sq)
  new_pis, new_means, new_sigmas = [
      tf.cast(x, tf.float32) for x in [new_pis, new_means, new_sigmas]]
  return new_pis, new_means, new_sigmas
//  [python]
def make_images_and_labels_tensors(examples_to_take):
    """Make tensors for loading images and labels from dataset."""

    with tf.name_scope("input"):
        dataset = dataset_utils.get_dataset(
            FLAGS.primary_dataset_name, FLAGS.split
        )
        dataset = dataset.filter(make_labeled_data_filter())

        # This is necessary for datasets that aren't shuffled on disk, such as
        # ImageNet.
        if FLAGS.split == "train":
            dataset = dataset.shuffle(FLAGS.shuffle_buffer_size, 0)

        # Optionally only use a certain fraction of the dataset.
        # This is used in at least 2 contexts:
        # 1. We don't evaluate on all training data sometimes for speed reasons.
        # 2. We may want smaller validation sets to see whether HPO still works.
        if examples_to_take != -1:
            dataset = dataset.take(examples_to_take)

        # Batch the results
        dataset = dataset.batch(FLAGS.batch_size)

        # Get the actual results from the iterator
        iterator = dataset.make_initializable_iterator()
        tf.add_to_collection(
            tf.GraphKeys.TABLE_INITIALIZERS, iterator.initializer
        )
        images, labels, _ = iterator.get_next()
        images = tf.to_float(images)

    return images, labels
//  [python]
def make_images_and_labels_tensors(examples_to_take):
    """Make tensors for loading images and labels from dataset."""

    with tf.name_scope("input"):
        dataset = dataset_utils.get_dataset(
            FLAGS.primary_dataset_name, FLAGS.split
        )
        dataset = dataset.filter(make_labeled_data_filter())

        # This is necessary for datasets that aren't shuffled on disk, such as
        # ImageNet.
        if FLAGS.split == "train":
            dataset = dataset.shuffle(FLAGS.shuffle_buffer_size, 0)

        # Optionally only use a certain fraction of the dataset.
        # This is used in at least 2 contexts:
        # 1. We don't evaluate on all training data sometimes for speed reasons.
        # 2. We may want smaller validation sets to see whether HPO still works.
        if examples_to_take != -1:
            dataset = dataset.take(examples_to_take)

        # Batch the results
        dataset = dataset.batch(FLAGS.batch_size)

        # Get the actual results from the iterator
        iterator = dataset.make_initializable_iterator()
        tf.add_to_collection(
            tf.GraphKeys.TABLE_INITIALIZERS, iterator.initializer
        )
        images, labels, _ = iterator.get_next()
        images = tf.cast(images, tf.float32)

    return images, labels
//  [python]
def moments_of_masked_time_series(time_series_tensor, broadcast_mask):
  """Compute mean and variance, accounting for a mask.

  Args:
    time_series_tensor: float `Tensor` time series of shape
      `concat([batch_shape, [num_timesteps]])`.
    broadcast_mask: bool `Tensor` of the same shape as `time_series`.
  Returns:
    mean: float `Tensor` of shape `batch_shape`.
    variance: float `Tensor` of shape `batch_shape`.
  """
  num_unmasked_entries = tf.cast(
      tf.reduce_sum(input_tensor=tf.cast(~broadcast_mask, tf.int32), axis=-1),
      time_series_tensor.dtype)

  # Manually compute mean and variance, excluding masked entries.
  mean = (
      tf.reduce_sum(
          input_tensor=tf.compat.v1.where(broadcast_mask,
                                          tf.zeros_like(time_series_tensor),
                                          time_series_tensor),
          axis=-1) / num_unmasked_entries)
  variance = (
      tf.reduce_sum(
          input_tensor=tf.compat.v1.where(
              broadcast_mask, tf.zeros_like(time_series_tensor),
              (time_series_tensor - mean[..., tf.newaxis])**2),
          axis=-1) / num_unmasked_entries)
  return mean, variance
//  [python]
def initial_value_of_masked_time_series(time_series_tensor, broadcast_mask):
  """Get the first unmasked entry of each time series in the batch.

  Args:
    time_series_tensor: float `Tensor` of shape [..., num_timesteps].
    broadcast_mask: bool `Tensor` of same shape as `time_series`.
  """

  num_timesteps = tf.shape(input=time_series_tensor)[-1]

  # Compute the index of the first unmasked entry for each series in the batch.
  unmasked_negindices = (
      tf.to_int32(~broadcast_mask) *
      tf.range(num_timesteps, 0, -1))
  first_unmasked_indices = num_timesteps - tf.reduce_max(
      input_tensor=unmasked_negindices, axis=-1)

  if first_unmasked_indices.shape.ndims is None:
    raise NotImplementedError(
        'Cannot compute initial values of a masked time series with'
        'dynamic rank.')  # `batch_gather` requires static rank

  # Extract the initial value for each series in the batch.
  return tf.squeeze(tf.compat.v1.batch_gather(
      params=time_series_tensor,
      indices=first_unmasked_indices[..., tf.newaxis]), axis=-1)
//  [python]
def initial_value_of_masked_time_series(time_series_tensor, broadcast_mask):
  """Get the first unmasked entry of each time series in the batch.

  Args:
    time_series_tensor: float `Tensor` of shape [..., num_timesteps].
    broadcast_mask: bool `Tensor` of same shape as `time_series`.
  """

  num_timesteps = tf.shape(input=time_series_tensor)[-1]

  # Compute the index of the first unmasked entry for each series in the batch.
  unmasked_negindices = (
      tf.cast(~broadcast_mask, tf.int32) *
      tf.range(num_timesteps, 0, -1))
  first_unmasked_indices = num_timesteps - tf.reduce_max(
      input_tensor=unmasked_negindices, axis=-1)

  if first_unmasked_indices.shape.ndims is None:
    raise NotImplementedError(
        'Cannot compute initial values of a masked time series with'
        'dynamic rank.')  # `batch_gather` requires static rank

  # Extract the initial value for each series in the batch.
  return tf.squeeze(tf.compat.v1.batch_gather(
      params=time_series_tensor,
      indices=first_unmasked_indices[..., tf.newaxis]), axis=-1)
//  [python]
def build_reward(self):
        with tf.name_scope('permutations'):
            # Reorder input % tour
            self.permutations = tf.stack(
                [
                    tf.tile(tf.expand_dims(tf.range(self.batch_size, dtype=tf.int32), 1), [1, self.max_length + 1]),
                    self.positions
                ],
                2
            )
            # Sort
            self.ordered_input_ = tf.gather_nd(self.input_, self.permutations)
            # [batch size, seq length + 1, features] to [features, seq length + 1, batch_size]
            self.ordered_input_ = tf.transpose(self.ordered_input_, [2, 1, 0])

            # Ordered coordinates
            ordered_x_ = self.ordered_input_[0]  # [seq length + 1, batch_size]
            # [batch_size, seq length]        delta_x**2
            delta_x2 = tf.transpose(tf.square(ordered_x_[1:] - ordered_x_[:-1]), [1, 0])
            ordered_y_ = self.ordered_input_[1]  # [seq length + 1, batch_size]
            # [batch_size, seq length]        delta_y**2
            delta_y2 = tf.transpose(tf.square(ordered_y_[1:] - ordered_y_[:-1]), [1, 0])

        with tf.name_scope('environment'):
            # Get tour length (euclidean distance)
            inter_city_distances = tf.sqrt(delta_x2 + delta_y2)  # sqrt(delta_x**2 + delta_y**2)
            self.distances = tf.reduce_sum(inter_city_distances, axis=1)  # [batch_size]

            self.reward = tf.to_float(self.distances)
//  [python]
def build_reward(self):
        with tf.name_scope('permutations'):
            # Reorder input % tour
            self.permutations = tf.stack(
                [
                    tf.tile(tf.expand_dims(tf.range(self.batch_size, dtype=tf.int32), 1), [1, self.max_length + 1]),
                    self.positions
                ],
                2
            )
            # Sort
            self.ordered_input_ = tf.gather_nd(self.input_, self.permutations)
            # [batch size, seq length + 1, features] to [features, seq length + 1, batch_size]
            self.ordered_input_ = tf.transpose(self.ordered_input_, [2, 1, 0])

            # Ordered coordinates
            ordered_x_ = self.ordered_input_[0]  # [seq length + 1, batch_size]
            # [batch_size, seq length]        delta_x**2
            delta_x2 = tf.transpose(tf.square(ordered_x_[1:] - ordered_x_[:-1]), [1, 0])
            ordered_y_ = self.ordered_input_[1]  # [seq length + 1, batch_size]
            # [batch_size, seq length]        delta_y**2
            delta_y2 = tf.transpose(tf.square(ordered_y_[1:] - ordered_y_[:-1]), [1, 0])

        with tf.name_scope('environment'):
            # Get tour length (euclidean distance)
            inter_city_distances = tf.sqrt(delta_x2 + delta_y2)  # sqrt(delta_x**2 + delta_y**2)
            self.distances = tf.reduce_sum(inter_city_distances, axis=1)  # [batch_size]

            self.reward = tf.cast(self.distances, tf.float32)
//  [python]
def get_padding(x, padding_value=0, dtype=tf.float32):
  """Return float tensor representing the padding values in x.

  Args:
    x: int tensor with any shape
    padding_value: int which represents padded values in input
    dtype: The dtype of the return value.

  Returns:
    float tensor with same shape as x containing values 0 or 1.
      0 -> non-padding, 1 -> padding
  """
  with tf.name_scope("padding"):
    return tf.cast(tf.equal(x, padding_value), dtype)
//  [python]
def get_position_encoding(length,
                          hidden_size,
                          min_timescale=1.0,
                          max_timescale=1.0e4):
  """Return positional encoding.

  Calculates the position encoding as a mix of sine and cosine functions with
  geometrically increasing wavelengths.
  Defined and formulized in Attention is All You Need, section 3.5.

  Args:
    length: Sequence length.
    hidden_size: Size of the
    min_timescale: Minimum scale that will be applied at each position
    max_timescale: Maximum scale that will be applied at each position

  Returns:
    Tensor with shape [length, hidden_size]
  """
  # We compute the positional encoding in float32 even if the model uses
  # float16, as many of the ops used, like log and exp, are numerically unstable
  # in float16.
  position = t.cast(t.range(length), t.float32)
  num_timescales = hidden_size // 2
  log_timescale_increment = (
      math.log(float(max_timescale) / float(min_timescale)) /
      (t.cast(num_timescales, t.float32) - 1))
  inv_timescales = min_timescale * t.exp(
      t.cast(t.range(num_timescales), t.float32) * -log_timescale_increment)
  scaled_time = t.expand_dims(position, 1) * t.expand_dims(inv_timescales, 0)
  signal = t.concat([t.sin(scaled_time), t.cos(scaled_time)], axis=1)
  return signal
//  [python]
def get_position_encoding(length,
                          hidden_size,
                          min_timescale=1.0,
                          max_timescale=1.0e4):
  """Return positional encoding.

  Calculates the position encoding as a mix of sine and cosine functions with
  geometrically increasing wavelengths.
  Defined and formulized in Attention is All You Need, section 3.5.

  Args:
    length: Sequence length.
    hidden_size: Size of the
    min_timescale: Minimum scale that will be applied at each position
    max_timescale: Maximum scale that will be applied at each position

  Returns:
    Tensor with shape [length, hidden_size]
  """
  # We compute the positional encoding in float32 even if the model uses
  # float16, as many of the ops used, like log and exp, are numerically unstable
  # in float16.
  position = tf.cast(tf.range(length), tf.float32)
  num_timescales = hidden_size // 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)
  return signal
//  [python]
def _combine_masks(mask1, mask2, dtype, how="and"):
  """Combines two masks.

  Use "and" if trying to combine two existing masks.
  Use "or" if trying to flip a few positions to "real".

  Args:
    mask1: tf.Tensor, input mask 1
    mask2: tf.Tensor, input mask 2
    dtype: tf.dtype
    how: Which logical operation should run.

  Returns:
    The combined input masks.

  """
  if how == "and":
    operator = tf.math.logical_and
  else:
    operator = tf.math.logical_or
  return tf.cast(operator(
      tf.cast(mask1, tf.bool),
      tf.cast(mask2, tf.bool)), dtype=dtype)
//  [python]
def get_num_items(field: Value, graph_shape: tf.Tensor) -> tf.Tensor:
  """Returns the number of items for each graph in a GraphTensor.

  Args:
    field: tensor with shape `[*graph_shape, num_items, *feature_shape]`,
      possibly ragged.
    graph_shape: rank-1 tensor with the graph shape.

  Returns:
    The number of items in `field` for each graph in its GraphTensor,
    as dense tensor with a shape `[*graph_shape]`.
  """
  if graph_shape.shape.rank != 1:
    raise ValueError(
        f'graph_shape must have rank=1, got {graph_shape.shape.rank}'
    )

  graph_rank = graph_shape.shape[0]
  if graph_rank is None:
    raise ValueError('graph_shape tensor must have fully defined shape')

  if field.shape.rank <= graph_rank:
    raise ValueError(
        f'Expected that value rank={field.shape.rank} is greater than'
        f' graph rank={graph_rank}'
    )

  def _get_num_items_flat(value, dim, dtype) -> tf.Tensor:
    """Retuns num items, as a scalar if uniform or a flat vector if not."""
    if is_dense_tensor(value):
      return tf.shape(value, out_type=dtype)[dim]
    elif is_ragged_tensor(value):
      if dim == 0:
        return tf.size(value.row_splits, out_type=dtype) - tf.constant(
            1, dtype=dtype
        )
      elif dim == 1:
        if value.uniform_row_length is not None:
          return tf.to_int32(value.uniform_row_length, dtype=dtype)
        else:
          return tf.to_int32(value.row_lengths(), dtype=dtype)
      else:
        return _get_num_items_flat(value.values, dim - 1, dtype=dtype)
    else:
      raise ValueError(f'Unsupported type {type(value).__name__}')

  num_items = _get_num_items_flat(field, graph_rank, dtype=graph_shape.dtype)
  if num_items.shape.rank == 0:
    return tf.fill(graph_shape, num_items)
  else:
    return tf.reshape(num_items, graph_shape)
//  [python]
def get_num_items(field: Value, graph_shape: tf.Tensor) -> tf.Tensor:
  """Returns the number of items for each graph in a GraphTensor.

  Args:
    field: tensor with shape `[*graph_shape, num_items, *feature_shape]`,
      possibly ragged.
    graph_shape: rank-1 tensor with the graph shape.

  Returns:
    The number of items in `field` for each graph in its GraphTensor,
    as dense tensor with a shape `[*graph_shape]`.
  """
  if graph_shape.shape.rank != 1:
    raise ValueError(
        f'graph_shape must have rank=1, got {graph_shape.shape.rank}'
    )

  graph_rank = graph_shape.shape[0]
  if graph_rank is None:
    raise ValueError('graph_shape tensor must have fully defined shape')

  if field.shape.rank <= graph_rank:
    raise ValueError(
        f'Expected that value rank={field.shape.rank} is greater than'
        f' graph rank={graph_rank}'
    )

  def _get_num_items_flat(value, dim, dtype) -> tf.Tensor:
    """Retuns num items, as a scalar if uniform or a flat vector if not."""
    if is_dense_tensor(value):
      return tf.shape(value, out_type=dtype)[dim]
    elif is_ragged_tensor(value):
      if dim == 0:
        return tf.size(value.row_splits, out_type=dtype) - tf.constant(
            1, dtype=dtype
        )
      elif dim == 1:
        if value.uniform_row_length is not None:
          return tf.cast(value.uniform_row_length, dtype=dtype)
        else:
          return tf.cast(value.row_lengths(), dtype=dtype)
      else:
        return _get_num_items_flat(value.values, dim - 1, dtype=dtype)
    else:
      raise ValueError(f'Unsupported type {type(value).__name__}')

  num_items = _get_num_items_flat(field, graph_rank, dtype=graph_shape.dtype)
  if num_items.shape.rank == 0:
    return tf.fill(graph_shape, num_items)
  else:
    return tf.reshape(num_items, graph_shape)
//  [python]
def select_fg_for_masks(class_targets, box_targets, boxes, proposal_to_label_map, max_num_fg=128):
    """Selects the fore ground objects for mask branch during training.

    Args:
    class_targets: a tensor of shape [batch_size, num_boxes]  representing the
      class label for each box.
    box_targets: a tensor with a shape of [batch_size, num_boxes, 4]. The tensor
      contains the ground truth pixel coordinates of the scaled images for each
      roi.
    boxes: A 3-D Tensor of shape [batch_size, num_boxes, 4]. Each row
      represents a box with [y1, x1, y2, x2] in un-normalized coordinates.
    proposal_to_label_map: a tensor with a shape of [batch_size, num_boxes].
      This tensor keeps the mapping between proposal to labels.
      proposal_to_label_map[i] means the index of the ground truth instance for
      the i-th proposal.
    max_num_fg: a integer represents the number of masks per image.
    Returns:
    class_targets, boxes, proposal_to_label_map, box_targets that have
    foreground objects.
    """

    # Masks are for positive (fg) objects only.
    # Reference: https://github.com/facebookresearch/Detectron/blob/master/detectron/roi_data/mask_rcnn.py
    batch_size = boxes.shape[0]
    _, fg_indices = tf.nn.top_k(tf.to_float(tf.greater(class_targets, 0)), k=max_num_fg)

    # Contructs indices for gather.
    indices = tf.reshape(
        fg_indices + tf.expand_dims(tf.range(batch_size) * tf.shape(input=class_targets)[1], 1),
        shape=[-1]
    )

    fg_class_targets = tf.reshape(
        tf.gather(tf.reshape(class_targets, [-1, 1]), indices),
        [batch_size, -1]
    )

    fg_box_targets = tf.reshape(
        tf.gather(tf.reshape(box_targets, [-1, 4]), indices),
        [batch_size, -1, 4]
    )

    fg_box_rois = tf.reshape(
        tf.gather(tf.reshape(boxes, [-1, 4]), indices), [batch_size, -1, 4]
    )

    fg_proposal_to_label_map = tf.reshape(
        tf.gather(tf.reshape(proposal_to_label_map, [-1, 1]), indices),
        [batch_size, -1]
    )

    return (fg_class_targets, fg_box_targets, fg_box_rois,
            fg_proposal_to_label_map)
//  [python]
def select_fg_for_masks(class_targets, box_targets, boxes, proposal_to_label_map, max_num_fg=128):
    """Selects the fore ground objects for mask branch during training.

    Args:
    class_targets: a tensor of shape [batch_size, num_boxes]  representing the
      class label for each box.
    box_targets: a tensor with a shape of [batch_size, num_boxes, 4]. The tensor
      contains the ground truth pixel coordinates of the scaled images for each
      roi.
    boxes: A 3-D Tensor of shape [batch_size, num_boxes, 4]. Each row
      represents a box with [y1, x1, y2, x2] in un-normalized coordinates.
    proposal_to_label_map: a tensor with a shape of [batch_size, num_boxes].
      This tensor keeps the mapping between proposal to labels.
      proposal_to_label_map[i] means the index of the ground truth instance for
      the i-th proposal.
    max_num_fg: a integer represents the number of masks per image.
    Returns:
    class_targets, boxes, proposal_to_label_map, box_targets that have
    foreground objects.
    """

    # Masks are for positive (fg) objects only.
    # Reference: https://github.com/facebookresearch/Detectron/blob/master/detectron/roi_data/mask_rcnn.py
    batch_size = boxes.shape[0]
    _, fg_indices = tf.nn.top_k(tf.cast(tf.greater(class_targets, 0), dtype=tf.float32), k=max_num_fg)

    # Contructs indices for gather.
    indices = tf.reshape(
        fg_indices + tf.expand_dims(tf.range(batch_size) * tf.shape(input=class_targets)[1], 1),
        shape=[-1]
    )

    fg_class_targets = tf.reshape(
        tf.gather(tf.reshape(class_targets, [-1, 1]), indices),
        [batch_size, -1]
    )

    fg_box_targets = tf.reshape(
        tf.gather(tf.reshape(box_targets, [-1, 4]), indices),
        [batch_size, -1, 4]
    )

    fg_box_rois = tf.reshape(
        tf.gather(tf.reshape(boxes, [-1, 4]), indices), [batch_size, -1, 4]
    )

    fg_proposal_to_label_map = tf.reshape(
        tf.gather(tf.reshape(proposal_to_label_map, [-1, 1]), indices),
        [batch_size, -1]
    )

    return (fg_class_targets, fg_box_targets, fg_box_rois,
            fg_proposal_to_label_map)
//  [python]
def finalize(beam_width, parent_ids, sequence_lengths, outputs, end_id, max_seq_len=None):
    maximum_lengths = tf.reduce_max(tf.reshape(
        sequence_lengths, [-1, beam_width]), axis=-1)
    
    if max_seq_len != None:
        array_shape = [max_seq_len, -1, beam_width]
    else:
        array_shape = [tf.reduce_max(maximum_lengths), -1, beam_width]
        
    step_ids = tf.reshape(outputs, array_shape)
    parent_ids = tf.reshape(parent_ids, array_shape)

    ids = tf.contrib.seq2seq.gather_tree(
        step_ids, parent_ids, maximum_lengths, end_id)

    ids = tf.transpose(ids, perm=[1, 2, 0])
    lengths = tf.not_equal(ids, end_id)
    lengths = tf.to_int32(lengths)
    lengths = tf.reduce_sum(lengths, axis=-1)
    return ids, lengths