jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def _build_tables(self, prior, precision, offset=None):
    # Prevent tensors from bouncing back and forth between host and GPU.

    with tf.device("/cpu:0"):
      # Variant of the base class where we use FOO
      precision = int(precision)
      offset = tf.cast(0 if offset is None else offset, prior.dtype)
      # Subclasses should have already caught this, but better be safe.
      assert not prior.event_shape.rank

      lower_tail = tfc.distributions.lower_tail(prior, self.tail_mass)
      upper_tail = tfc.distributions.upper_tail(prior, self.tail_mass)
      # Integers such that:
      # minima + offset < lower_tail
      # maxima + offset > upper_tail
      minima = tf.cast(tf.math.floor(lower_tail - offset), tf.int32)
      maxima = tf.cast(tf.math.ceil(upper_tail - offset), tf.int32)

      # PMF starting positions and lengths.
      pmf_start = tf.cast(minima, prior.dtype) + offset
      pmf_length = maxima - minima + 1

      # Sample the densities in the computed ranges, possibly computing more
      # samples than necessary at the upper end.
      max_length = tf.math.reduce_max(pmf_length)
      if tf.executing_eagerly() and max_length > 2048:
        logging.warning(
            "Very wide PMF with %d elements may lead to out of memory issues. "
            "Consider priors with smaller variance, or increasing `tail_mass` "
            "parameter.", int(max_length))
      samples = tf.range(tf.cast(max_length, prior.dtype), dtype=prior.dtype)
      samples = tf.reshape(samples, [-1] + pmf_length.shape.rank * [1])
      samples += pmf_start

      # TODO(b/251395094): Move this into TFC.
      # NOTE: This line is the main difference to the base class: Here we use
      # exp(log_prob(x)), which includes the Laplacian tail mass.
      pmf = tf.math.exp(self._log_prob(prior, samples))
      pmf_shape = tf.shape(pmf)[1:]
      num_pmfs = tf.reduce_prod(pmf_shape)

      # Collapse batch dimensions of distribution.
      pmf = tf.reshape(pmf, [max_length, num_pmfs])
      pmf = tf.transpose(pmf)

      logging.info("PMF shape: %s", pmf.shape)

      pmf_length = tf.broadcast_to(pmf_length, pmf_shape)
      pmf_length = tf.reshape(pmf_length, [num_pmfs])
      cdf_offset = tf.broadcast_to(minima, pmf_shape)
      cdf_offset = tf.reshape(cdf_offset, [num_pmfs])

      cdf = get_cdf(pmf, pmf_length, int(num_pmfs.numpy()), precision)

    return cdf, cdf_offset
//  [python]
def get_saliency(attn_map, shape=(128, 128), method='nearest'):
    bs = attn_map.shape[0]
    attn_map = tf.reshape(attn_map, [bs, -1, 256])
    non_zeros = tf.math.count_nonzero(attn_map, axis=1)
    weights = tf.math.reduce_sum(attn_map, axis=1) / tf.to_float(non_zeros)
    weights = tf.reshape(weights, [bs, 16, 16, 1])
    weights = tf.image.resize(weights, shape, method=method)
    return weights
//  [python]
def get_saliency(attn_map, shape=(128, 128), method='nearest'):
    bs = attn_map.shape[0]
    attn_map = tf.reshape(attn_map, [bs, -1, 256])
    non_zeros = tf.math.count_nonzero(attn_map, axis=1)
    weights = tf.math.reduce_sum(attn_map, axis=1) / tf.cast(non_zeros, tf.float32)
    weights = tf.reshape(weights, [bs, 16, 16, 1])
    weights = tf.image.resize(weights, shape, method=method)
    return weights
//  [python]
def get_sampling_indices(sampling_seed, class_counts, labels, num_samples):
  """Function to generates sampling indices of the given labels.

  Args:
    sampling_seed: seed used for the sampling process.
    class_counts: instances count of every class in the dataset. If None,
      uniform sampling is applied.
    labels: ground-truth labels from the target dataset.
    num_samples: number of labels to sample from the labels.

  Returns:
    indices: indices for sampling the labels.

  """
  np.random.seed(sampling_seed)
  num_samples = min(num_samples, labels.shape[0])  # To avoid sampling errors

  if not class_counts:  # Uniform sampling
    weights_labels = None
  else:  # Class balanced sampling
    weights_labels = [1/class_counts[int(x)] for x in labels]
    weights_labels = weights_labels / np.sum(weights_labels)

  indices = np.random.choice(
      np.arange(labels.shape[0]),
      num_samples,
      replace=False,
      p=weights_labels)
  return t.cast(indices, t.int32)
//  [python]
def get_sampling_indices(sampling_seed, class_counts, labels, num_samples):
  """Function to generates sampling indices of the given labels.

  Args:
    sampling_seed: seed used for the sampling process.
    class_counts: instances count of every class in the dataset. If None,
      uniform sampling is applied.
    labels: ground-truth labels from the target dataset.
    num_samples: number of labels to sample from the labels.

  Returns:
    indices: indices for sampling the labels.

  """
  np.random.seed(sampling_seed)
  num_samples = min(num_samples, labels.shape[0])  # To avoid sampling errors

  if not class_counts:  # Uniform sampling
    weights_labels = None
  else:  # Class balanced sampling
    weights_labels = [1/class_counts[int(x)] for x in labels]
    weights_labels = weights_labels / np.sum(weights_labels)

  indices = np.random.choice(
      np.arange(labels.shape[0]),
      num_samples,
      replace=False,
      p=weights_labels)
  return tf.cast(indices, tf.int32)
//  [python]
def switch(condition, then_expression, else_expression):
    """
    Switches between two operations depending on a scalar value (int or bool).
    Note that both `then_expression` and `else_expression`
    should be symbolic tensors of the *same shape*.

    :param condition: (TensorFlow Tensor) scalar tensor.
    :param then_expression: (TensorFlow Operation)
    :param else_expression: (TensorFlow Operation)
    :return: (TensorFlow Operation) the switch output
    """
    x_shape = copy.copy(then_expression.get_shape())
    out_tensor = tf.cond(t.cast(condition, 'bool'),
                         lambda: then_expression,
                         lambda: else_expression)
    out_tensor.set_shape(x_shape)
    return out_tensor
//  [python]
def switch(condition, then_expression, else_expression):
    """
    Switches between two operations depending on a scalar value (int or bool).
    Note that both `then_expression` and `else_expression`
    should be symbolic tensors of the *same shape*.

    :param condition: (TensorFlow Tensor) scalar tensor.
    :param then_expression: (TensorFlow Operation)
    :param else_expression: (TensorFlow Operation)
    :return: (TensorFlow Operation) the switch output
    """
    x_shape = copy.copy(then_expression.get_shape())
    out_tensor = tf.cond(tf.cast(condition, 'bool'),
                         lambda: then_expression,
                         lambda: else_expression)
    out_tensor.set_shape(x_shape)
    return out_tensor
//  [python]
preprocess_img = preprocess(sample_img)
print(f"Preprocessed image shape: {preprocess_img.shape}")
print(
    f"Preprocessed image pixel range: ({preprocess_img.numpy().min()},",
    f"{preprocess_img.numpy().max()})",
)

batched_img = tf.expand_dims(preprocess_img, axis=0)
batched_img = tf.to_float(batched_img)
//  [python]
preprocess_img = preprocess(sample_img)
print(f"Preprocessed image shape: {preprocess_img.shape}")
print(
    f"Preprocessed image pixel range: ({preprocess_img.numpy().min()},",
    f"{preprocess_img.numpy().max()})",
)

batched_img = tf.expand_dims(preprocess_img, axis=0)
batched_img = tf.cast(batched_img, tf.float32)
//  [python]
def call(self, inputs, global_step=None, training=None):
    # Computes random features.
    gp_inputs = inputs
    if self.normalize_input:
      gp_inputs = self._input_norm_layer(gp_inputs)
    elif self.use_custom_random_features:
      # Supports lengthscale for custom random feature layer by directly
      # rescaling the input.
      gp_input_scale = t.cast(self.gp_input_scale, inputs.dtype)
      gp_inputs = gp_inputs * gp_input_scale

    gp_feature = self._random_feature(gp_inputs)

    if self.scale_random_features:
      # Scale random feature by 2. / sqrt(num_inducing) following [1].
      # When using GP layer as the output layer of a nerual network,
      # it is recommended to turn this scaling off to prevent it from changing
      # the learning rate to the hidden layers.
      gp_feature_scale = t.cast(self.gp_feature_scale, inputs.dtype)
      gp_feature = gp_feature * gp_feature_scale

    # Computes posterior center (i.e., MAP estimate) and variance.
    gp_output = self._gp_output_layer(gp_feature) + self._gp_output_bias

    if self.return_gp_cov:
      gp_covmat = self._gp_cov_layer(gp_feature, gp_output, training)

    # Assembles model output.
    model_output = [gp_output,]
    if self.return_gp_cov:
      model_output.append(gp_covmat)
    if self.return_random_features:
      model_output.append(gp_feature)

    return model_output
//  [python]
def call(self, inputs, global_step=None, training=None):
    # Computes random features.
    gp_inputs = inputs
    if self.normalize_input:
      gp_inputs = self._input_norm_layer(gp_inputs)
    elif self.use_custom_random_features:
      # Supports lengthscale for custom random feature layer by directly
      # rescaling the input.
      gp_input_scale = tf.cast(self.gp_input_scale, inputs.dtype)
      gp_inputs = gp_inputs * gp_input_scale

    gp_feature = self._random_feature(gp_inputs)

    if self.scale_random_features:
      # Scale random feature by 2. / sqrt(num_inducing) following [1].
      # When using GP layer as the output layer of a nerual network,
      # it is recommended to turn this scaling off to prevent it from changing
      # the learning rate to the hidden layers.
      gp_feature_scale = tf.cast(self.gp_feature_scale, inputs.dtype)
      gp_feature = gp_feature * gp_feature_scale

    # Computes posterior center (i.e., MAP estimate) and variance.
    gp_output = self._gp_output_layer(gp_feature) + self._gp_output_bias

    if self.return_gp_cov:
      gp_covmat = self._gp_cov_layer(gp_feature, gp_output, training)

    # Assembles model output.
    model_output = [gp_output,]
    if self.return_gp_cov:
      model_output.append(gp_covmat)
    if self.return_random_features:
      model_output.append(gp_feature)

    return model_output
//  [python]
def make_precision_matrix_update_op(self,
                                      gp_feature,
                                      logits,
                                      precision_matrix):
    """Defines update op for the precision matrix of feature weights."""
    if self.likelihood != 'gaussian':
      if logits is None:
        raise ValueError(
            f'"logits" cannot be None when likelihood={self.likelihood}')

      if logits.shape[-1] != 1:
        raise ValueError(
            f'likelihood={self.likelihood} only support univariate logits.'
            f'Got logits dimension: {logits.shape[-1]}')

    batch_size = tf.shape(gp_feature)[0]
    batch_size = tf.to_float(batch_size)

    # Computes batch-specific normalized precision matrix.
    if self.likelihood == 'binary_logistic':
      prob = tf.sigmoid(logits)
      prob_multiplier = prob * (1. - prob)
    elif self.likelihood == 'poisson':
      prob_multiplier = tf.exp(logits)
    else:
      prob_multiplier = 1.

    gp_feature_adjusted = tf.sqrt(prob_multiplier) * gp_feature
    precision_matrix_minibatch = tf.matmul(
        gp_feature_adjusted, gp_feature_adjusted, transpose_a=True)

    # Updates the population-wise precision matrix.
    if self.momentum > 0:
      # Use moving-average updates to accumulate batch-specific precision
      # matrices.
      precision_matrix_minibatch = precision_matrix_minibatch / batch_size
      precision_matrix_new = (
          self.momentum * precision_matrix +
          (1. - self.momentum) * precision_matrix_minibatch)
    else:
      # Compute exact population-wise covariance without momentum.
      # If use this option, make sure to pass through data only once.
      precision_matrix_new = precision_matrix + precision_matrix_minibatch

    # Returns the update op.
    return precision_matrix.assign(precision_matrix_new)
//  [python]
def make_precision_matrix_update_op(self,
                                      gp_feature,
                                      logits,
                                      precision_matrix):
    """Defines update op for the precision matrix of feature weights."""
    if self.likelihood != 'gaussian':
      if logits is None:
        raise ValueError(
            f'"logits" cannot be None when likelihood={self.likelihood}')

      if logits.shape[-1] != 1:
        raise ValueError(
            f'likelihood={self.likelihood} only support univariate logits.'
            f'Got logits dimension: {logits.shape[-1]}')

    batch_size = tf.shape(gp_feature)[0]
    batch_size = tf.cast(batch_size, dtype=gp_feature.dtype)

    # Computes batch-specific normalized precision matrix.
    if self.likelihood == 'binary_logistic':
      prob = tf.sigmoid(logits)
      prob_multiplier = prob * (1. - prob)
    elif self.likelihood == 'poisson':
      prob_multiplier = tf.exp(logits)
    else:
      prob_multiplier = 1.

    gp_feature_adjusted = tf.sqrt(prob_multiplier) * gp_feature
    precision_matrix_minibatch = tf.matmul(
        gp_feature_adjusted, gp_feature_adjusted, transpose_a=True)

    # Updates the population-wise precision matrix.
    if self.momentum > 0:
      # Use moving-average updates to accumulate batch-specific precision
      # matrices.
      precision_matrix_minibatch = precision_matrix_minibatch / batch_size
      precision_matrix_new = (
          self.momentum * precision_matrix +
          (1. - self.momentum) * precision_matrix_minibatch)
    else:
      # Compute exact population-wise covariance without momentum.
      # If use this option, make sure to pass through data only once.
      precision_matrix_new = precision_matrix + precision_matrix_minibatch

    # Returns the update op.
    return precision_matrix.assign(precision_matrix_new)
//  [python]
def _get_lm_mask(
      self, inputs_mask, branch1):
    """Dynamically selects tokens to be part of the Cloze task."""
    inputs_shape = tf.shape(inputs_mask)
    batched = len(inputs_shape) == 2
    batch_size = inputs_shape[0] if batched else None
    padded_len = inputs_shape[1] if batched else inputs_shape[0]
    # Select a proportion self._masked_lm_prob of the tokens in each sequence
    # uniformly at random...
    if self._deterministic_n_lm_tokens:
      # tf.function is finnicky with types...
      inputs_mask = tf.to_float(inputs_mask)
      masked_lm_prob = tf.to_float(self._masked_lm_prob)
      padded_len = tf.to_float(padded_len)

      n_normal_tokens_per_seq = tf.reduce_sum(inputs_mask, -1)
      max_lm_tokens_per_seq = tf.cast(masked_lm_prob * padded_len, tf.int32)
      lm_tokens_per_seq = tf.maximum(
          1, tf.cast(masked_lm_prob * n_normal_tokens_per_seq, tf.int32))

      indices = lm_tokens_per_seq - 1
      if batched:
        indices = tf.stack((tf.range(batch_size), indices), axis=-1)
      branch1 *= inputs_mask
      values, _ = math.top_k(branch1, k=max_lm_tokens_per_seq)
      cutoffs = (tf.gather(values, indices) if not batched
                 else tf.gather_nd(values, indices)[:, None])

      return branch1 >= cutoffs
    # ... otherwise, select them via i.i.d. coin tosses with success probability
    # equal to self._masked_lm_prob
    return tf.logical_and(inputs_mask, branch1)
//  [python]
def _get_lm_mask(
      self, inputs_mask, branch1):
    """Dynamically selects tokens to be part of the Cloze task."""
    inputs_shape = tf.shape(inputs_mask)
    batched = len(inputs_shape) == 2
    batch_size = inputs_shape[0] if batched else None
    padded_len = inputs_shape[1] if batched else inputs_shape[0]
    # Select a proportion self._masked_lm_prob of the tokens in each sequence
    # uniformly at random...
    if self._deterministic_n_lm_tokens:
      # tf.function is finnicky with types...
      inputs_mask = tf.cast(inputs_mask, tf.float32)
      masked_lm_prob = tf.cast(self._masked_lm_prob, tf.float32)
      padded_len = tf.cast(padded_len, tf.float32)

      n_normal_tokens_per_seq = tf.reduce_sum(inputs_mask, -1)
      max_lm_tokens_per_seq = tf.cast(masked_lm_prob * padded_len, tf.int32)
      lm_tokens_per_seq = tf.maximum(
          1, tf.cast(masked_lm_prob * n_normal_tokens_per_seq, tf.int32))

      indices = lm_tokens_per_seq - 1
      if batched:
        indices = tf.stack((tf.range(batch_size), indices), axis=-1)
      branch1 *= inputs_mask
      values, _ = math.top_k(branch1, k=max_lm_tokens_per_seq)
      cutoffs = (tf.gather(values, indices) if not batched
                 else tf.gather_nd(values, indices)[:, None])

      return branch1 >= cutoffs
    # ... otherwise, select them via i.i.d. coin tosses with success probability
    # equal to self._masked_lm_prob
    return tf.logical_and(inputs_mask, branch1)
//  [python]
def call(self, inputs):
    """Dynamically masks a batch of tokens for BERT-like pretraining.

    Args:
      inputs: A tf.Tensor<int32>[batch, len] of sequences.

    Returns:
      A 3-tuple of tf.Tensor<int32>[batch, len] representing the masked inputs,
      the target output (equal to the unmasked input) and the mask.
    """
    inputs_shape = tf.shape(inputs)
    # inputs_mask[i][j] = True iff token of sequence i at position j is a normal
    # token, i.e. neither padding nor <CLS> / <EOS>
    inputs_mask = self._vocab.padding_mask(inputs)
    inputs_mask = tf.math.logical_and(
        inputs_mask, self._vocab.special_token_mask(inputs))

    # Sample Bernoulli-distributed tensors that will decide the "branching" on a
    # token by token basis...
    branch1 = self._branch1_sampler.sample(inputs_shape)
    branch2 = self._branch2_sampler.sample(inputs_shape)
    branch3 = self._branch3_sampler.sample(inputs_shape)
    # ...as well as the tokens resampled uniformly at random
    resampled_inputs = self._resample_sampler.sample(inputs_shape)

    # Mask indicating which tokens form the Cloze task in this call
    lm_mask = self._get_lm_mask(inputs_mask, branch1)

    # Build the final input
    mask_token_cond = tf.logical_and(lm_mask, branch2)
    resample_token_cond = tf.logical_and(
        tf.logical_and(lm_mask, tf.logical_not(branch2)), branch3)

    masked_inputs = tf.where(mask_token_cond,
                             t.cast(self._vocab.mask_code, inputs.dtype),
                             inputs)
    masked_inputs = tf.where(resample_token_cond,
                             t.cast(resampled_inputs, masked_inputs.dtype),
                             masked_inputs)

    return masked_inputs, inputs, t.cast(lm_mask, tf.float32)
//  [python]
def call(self, inputs):
    """Dynamically masks a batch of tokens for BERT-like pretraining.

    Args:
      inputs: A tf.Tensor<int32>[batch, len] of sequences.

    Returns:
      A 3-tuple of tf.Tensor<int32>[batch, len] representing the masked inputs,
      the target output (equal to the unmasked input) and the mask.
    """
    inputs_shape = tf.shape(inputs)
    # inputs_mask[i][j] = True iff token of sequence i at position j is a normal
    # token, i.e. neither padding nor <CLS> / <EOS>
    inputs_mask = self._vocab.padding_mask(inputs)
    inputs_mask = tf.math.logical_and(
        inputs_mask, self._vocab.special_token_mask(inputs))

    # Sample Bernoulli-distributed tensors that will decide the "branching" on a
    # token by token basis...
    branch1 = self._branch1_sampler.sample(inputs_shape)
    branch2 = self._branch2_sampler.sample(inputs_shape)
    branch3 = self._branch3_sampler.sample(inputs_shape)
    # ...as well as the tokens resampled uniformly at random
    resampled_inputs = self._resample_sampler.sample(inputs_shape)

    # Mask indicating which tokens form the Cloze task in this call
    lm_mask = self._get_lm_mask(inputs_mask, branch1)

    # Build the final input
    mask_token_cond = tf.logical_and(lm_mask, branch2)
    resample_token_cond = tf.logical_and(
        tf.logical_and(lm_mask, tf.logical_not(branch2)), branch3)

    masked_inputs = tf.where(mask_token_cond,
                             tf.cast(self._vocab.mask_code, inputs.dtype),
                             inputs)
    masked_inputs = tf.where(resample_token_cond,
                             tf.cast(resampled_inputs, masked_inputs.dtype),
                             masked_inputs)

    return masked_inputs, inputs, tf.cast(lm_mask, tf.float32)
//  [python]
def prediction(self):
        # input features
        self.dnn_parent_scope = 'dnn'
        with tf.variable_scope(self.dnn_parent_scope):
            with tf.variable_scope("input_from_feature_columns",
                                   partitioner=self.input_layer_partitioner,
                                   reuse=tf.AUTO_REUSE) as dnn_inputs_scope:
                net = tf.feature_column.input_layer(
                    features=self.feature, feature_columns=self.deep_column)

                add_layer_summary(net, dnn_inputs_scope.name)

            if self.bf16:
                net = tf.cast(net, dtype=tf.bfloat16)
                with tf.variable_scope(
                        'dnn_layers',
                        partitioner=self.dense_layer_partitioner,
                        reuse=tf.AUTO_REUSE).keep_weights():
                    net = self.dnn(net, self.dnn_hidden_units, "hiddenlayer")

                    with tf.variable_scope(
                            "logits",
                            values=(net, )).keep_weights() as dnn_logits_scope:
                        dnn_logits = tf.layers.dense(net,
                                                     units=1,
                                                     activation=None,
                                                     name=dnn_logits_scope)
                    add_layer_summary(dnn_logits, dnn_logits_scope.name)
                dnn_logits = tf.cast(dnn_logits, dtype=tf.float32)

            else:
                with tf.variable_scope(
                        'dnn_layers',
                        partitioner=self.dense_layer_partitioner,
                        reuse=tf.AUTO_REUSE):
                    net = self.dnn(net, self.dnn_hidden_units, "hiddenlayer")

                with tf.variable_scope("logits",
                                       values=(net, )) as dnn_logits_scope:
                    dnn_logits = tf.layers.dense(net,
                                                 units=1,
                                                 activation=None,
                                                 name=dnn_logits_scope)
                add_layer_summary(dnn_logits, dnn_logits_scope.name)

        self.linear_parent_scope = 'linear'
        with tf.variable_scope(
                self.linear_parent_scope,
                partitioner=self.input_layer_partitioner) as scope:
            linear_logits = tf.feature_column.linear_model(
                units=1,
                features=self.feature,
                feature_columns=self.wide_column,
                sparse_combiner='sum',
                weight_collections=None,
                trainable=True)

            add_layer_summary(linear_logits, scope.name)

        self.logits = tf.add_n([dnn_logits, linear_logits])
        predict = tf.math.sigmoid(self.logits)

        return predict
//  [python]
def shape_similarity_loss(logit_1, logit_2, logit_3, cube_params_1,
    cube_params_2, cube_params_3, node_position, n_part_1, n_part_2):
  with tf.name_scope('shape_similarity_loss'):
    _, _, relation_12 = cube_coverage_loss(cube_params_1, cube_params_2, n_part_1,
        node_position) # [bs, n_part_1]
    _, _, relation_23 = cube_coverage_loss(cube_params_2, cube_params_3, n_part_2,
        node_position) # [bs, n_part_2]
    coverage_loss_1, point_count_1 = coverage_split_loss(cube_params_1, node_position) # [bs, n_part_1]
    coverage_loss_2, point_count_2 = coverage_split_loss(cube_params_2, node_position) # [bs, n_part_2]
    coverage_loss_3, point_count_3 = coverage_split_loss(cube_params_3, node_position) # [bs, n_part_3]
    total_coverage_loss = tf.concat([coverage_loss_1, coverage_loss_2, coverage_loss_3], axis=1) # [bs, n1+n2+n3]
    point_count = tf.cast(tf.concat([point_count_1, point_count_2, point_count_3], axis=1), tf.float32) # [bs, n1+n2+n3]
    consistency_loss_1 = consistency_split_loss(cube_params_1, node_position, num_sample=26) # [bs, n_part_1]
    consistency_loss_2 = consistency_split_loss(cube_params_2, node_position, num_sample=26) # [bs, n_part_2]
    consistency_loss_3 = consistency_split_loss(cube_params_3, node_position, num_sample=26) # [bs, n_part_3]
    total_consistency_loss = tf.concat([consistency_loss_1, consistency_loss_2, consistency_loss_3], axis=1) # [bs, n1+n2+n3]

    def compute_chamfer_distance(logit):
      mean_coverage_loss = tf.reduce_sum(total_coverage_loss*logit, axis=1, keepdims=True)/tf.reduce_sum(logit*point_count, axis=1, keepdims=True) # [bs, 1]
      mean_consistency_loss = tf.reduce_sum(total_consistency_loss*logit, axis=1, keepdims=True)/tf.reduce_sum(logit, axis=1, keepdims=True) # [bs, 1]
      mean_chamfer_distance = mean_coverage_loss + mean_consistency_loss # [bs, 1]
      return mean_chamfer_distance

    logit = tf.concat([logit_1, logit_2, logit_3], axis=1) # [bs, n1+n2+n3]
    chamfer_distance = compute_chamfer_distance(logit) # [bs, 1]
    max_logit_1 = tf.zeros_like(logit_1)
    max_logit_2 = tf.zeros_like(logit_2)
    max_logit_3 = tf.ones_like(logit_3)
    max_logit = tf.concat([max_logit_1, max_logit_2, max_logit_3], axis=1) # [bs, n1+n2+n3]
    max_chamfer_distance = compute_chamfer_distance(max_logit) # [bs, 1]
    min_logit_1 = tf.ones_like(logit_1)
    min_logit_2 = tf.zeros_like(logit_2)
    min_logit_3 = tf.zeros_like(logit_3)
    min_logit = tf.concat([min_logit_1, min_logit_2, min_logit_3], axis=1) # [bs, n1+n2+n3]
    min_chamfer_distance = compute_chamfer_distance(min_logit) # [bs, 1]
    normalized_chamfer_distance = 2 * chamfer_distance / (max_chamfer_distance + min_chamfer_distance) # [bs, 1]
    loss = tf.reduce_mean(normalized_chamfer_distance) # [1]
  return loss, relation_12, relation_23
//  [python]
def read_tensor_from_image_file(image_file, input_height=299, input_width=299,
				input_mean=0, input_std=255):
    
  float_caster = tf.cast(image_file, tf.float32)
  dims_expander = tf.expand_dims(float_caster, 0);
  resized = tf.image.resize_bilinear(dims_expander, [input_height, input_width])
  normalized = tf.divide(tf.subtract(resized, [input_mean]), [input_std])
  sess = tf.Session()
  result = sess.run(normalized)

  return result
//  [python]
def _mp_hit_point_2_cell_ind(self, rot_gcs_2_mp, cm_center, cm_size,
                                 cm_cell_size, num_cells, hit_point):
        r"""
        Computes the indices of the cells to which points ``hit_point`` on the
        measurement plane belongs.

        Input
        ------
        rot_gcs_2_mp : [3, 3], tf.float
            Rotation matrix for going from the measurement plane LCS to the GCS

        cm_center : [3], tf.float
            Center of the coverage map

        cm_size : [2], tf.float
            Size of the coverage map

        cm_cell_size : [2], tf.float
            Size of of the cells of the ceverage map

        num_cells : [2], tf.int
            Number of cells in the coverage map

        hit_point : [...,3]
            Intersection points

        Output
        -------
        cell_ind : [..., 2], tf.int
            Indices of the cells
        """

        # Expand for broadcasting
        # [..., 3, 3]
        rot_gcs_2_mp = expand_to_rank(rot_gcs_2_mp, tf.rank(hit_point)+1, 0)
        # [..., 3]
        cm_center = expand_to_rank(cm_center, tf.rank(hit_point), 0)

        # Coverage map cells' indices
        # Coordinates of the hit point in the coverage map LCS
        # [..., 3]
        hit_point = tf.linalg.matvec(rot_gcs_2_mp, hit_point - cm_center)

        # In the local coordinate system of the coverage map, z should be 0
        # as the coverage map is in XY

        # x
        # [...]
        cell_x = hit_point[...,0] + cm_size[0]*0.5
        cell_x = tf.floor(cell_x/cm_cell_size[0])
        cell_x = tf.where(tf.less(cell_x, num_cells[0]), cell_x, num_cells[0])
        cell_x = tf.where(tf.greater_equal(cell_x, 0), cell_x, num_cells[0])

        # y
        # [...]
        cell_y = hit_point[...,1] + cm_size[1]*0.5
        cell_y = tf.floor(cell_y/cm_cell_size[1])
        cell_y = tf.where(tf.less(cell_y, num_cells[1]), cell_y, num_cells[1])
        cell_y = tf.where(tf.greater_equal(cell_y, 0), cell_y, num_cells[1])

        # [..., 2]
        cell_ind = tf.stack([cell_y, cell_x], axis=-1)

        return cell_ind
//  [python]
def _mp_hit_point_2_cell_ind(self, rot_gcs_2_mp, cm_center, cm_size,
                                 cm_cell_size, num_cells, hit_point):
        r"""
        Computes the indices of the cells to which points ``hit_point`` on the
        measurement plane belongs.

        Input
        ------
        rot_gcs_2_mp : [3, 3], tf.float
            Rotation matrix for going from the measurement plane LCS to the GCS

        cm_center : [3], tf.float
            Center of the coverage map

        cm_size : [2], tf.float
            Size of the coverage map

        cm_cell_size : [2], tf.float
            Size of of the cells of the ceverage map

        num_cells : [2], tf.int
            Number of cells in the coverage map

        hit_point : [...,3]
            Intersection points

        Output
        -------
        cell_ind : [..., 2], tf.int
            Indices of the cells
        """

        # Expand for broadcasting
        # [..., 3, 3]
        rot_gcs_2_mp = expand_to_rank(rot_gcs_2_mp, tf.rank(hit_point)+1, 0)
        # [..., 3]
        cm_center = expand_to_rank(cm_center, tf.rank(hit_point), 0)

        # Coverage map cells' indices
        # Coordinates of the hit point in the coverage map LCS
        # [..., 3]
        hit_point = tf.linalg.matvec(rot_gcs_2_mp, hit_point - cm_center)

        # In the local coordinate system of the coverage map, z should be 0
        # as the coverage map is in XY

        # x
        # [...]
        cell_x = hit_point[...,0] + cm_size[0]*0.5
        cell_x = tf.cast(tf.math.floor(cell_x/cm_cell_size[0]), tf.int32)
        cell_x = tf.where(tf.less(cell_x, num_cells[0]), cell_x, num_cells[0])
        cell_x = tf.where(tf.greater_equal(cell_x, 0), cell_x, num_cells[0])

        # y
        # [...]
        cell_y = hit_point[...,1] + cm_size[1]*0.5
        cell_y = tf.cast(tf.math.floor(cell_y/cm_cell_size[1]), tf.int32)
        cell_y = tf.where(tf.less(cell_y, num_cells[1]), cell_y, num_cells[1])
        cell_y = tf.where(tf.greater_equal(cell_y, 0), cell_y, num_cells[1])

        # [..., 2]
        cell_ind = tf.stack([cell_y, cell_x], axis=-1)

        return cell_ind
//  [python]
def call(self, input_shape: tf.TensorShape, past_key_values_length: int = 0):
        """Input is expected to be of size [bsz x seqlen]."""
        seq_len = input_shape[1]
        position_ids = tf.range(seq_len, delta=1, name="range")
        position_ids += past_key_values_length

        return super().call(tf.to_int32(position_ids))
//  [python]
def call(self, input_shape: tf.TensorShape, past_key_values_length: int = 0):
        """Input is expected to be of size [bsz x seqlen]."""
        seq_len = input_shape[1]
        position_ids = tf.range(seq_len, delta=1, name="range")
        position_ids += past_key_values_length

        return super().call(tf.cast(position_ids, dtype=tf.int32))
//  [python]
def _discard_obstructing_wedges(self, candidate_wedges, sources_positions):
        r"""
        Discard wedges for which the source is "inside" the wedge

        Input
        ------
        candidate_wedges : [num_candidate_wedges], int
            Candidate wedges.
            Entries correspond to wedges indices.

        sources_positions : [num_tx, 3], tf.float
            Coordinates of the sources.

        Output
        -------
        diff_mask : [num_tx, num_candidate_wedges], tf.bool
            Mask set to False for invalid wedges

        diff_wedges_ind : [num_candidate_wedges], tf.int
            Indices of the wedges that interacted with the diffracted paths
        """

        epsilon = t.cast(SolverBase.EPSILON, self._rdtype)

        # [num_candidate_wedges, 3]
        origins = tf.gather(self._wedges_origin, candidate_wedges)

        # Expand to broadcast with sources/targets and 0/n faces
        # [1, num_candidate_wedges, 1, 3]
        origins = tf.expand_dims(origins, axis=0)
        origins = tf.expand_dims(origins, axis=2)

        # Normals
        # [num_candidate_wedges, 2, 3]
        # [:,0,:] : 0-face
        # [:,1,:] : n-face
        normals = tf.gather(self._wedges_normals, candidate_wedges)
        # Expand to broadcast with the sources or targets
        # [1, num_candidate_wedges, 2, 3]
        normals = tf.expand_dims(normals, axis=0)

        # Expand to broadcast with candidate and 0/n faces wedges
        # [num_tx, 1, 1, 3]
        sources_positions = expand_to_rank(sources_positions, 4, 1)
        # Sources vectors
        # [num_tx, num_candidate_wedges, 1, 3]
        u_t = sources_positions - origins

        # [num_tx, num_candidate_wedges, 2]
        mask = dot(u_t, normals)
        mask = tf.greater(mask, tf.fill(tf.shape(mask), epsilon))
        # [num_tx, num_candidate_wedges]
        mask = tf.reduce_any(mask, axis=2)

        # Discard wedges with no valid link
        # [num_candidate_wedges]
        valid_wedges = tf.where(tf.reduce_any(mask, axis=0))[:,0]
        # [num_tx, num_candidate_wedges]
        mask = tf.gather(mask, valid_wedges, axis=1)
        # [num_candidate_wedges]
        diff_wedges_ind = tf.gather(candidate_wedges, valid_wedges, axis=0)

        return mask, diff_wedges_ind
//  [python]
def _discard_obstructing_wedges(self, candidate_wedges, sources_positions):
        r"""
        Discard wedges for which the source is "inside" the wedge

        Input
        ------
        candidate_wedges : [num_candidate_wedges], int
            Candidate wedges.
            Entries correspond to wedges indices.

        sources_positions : [num_tx, 3], tf.float
            Coordinates of the sources.

        Output
        -------
        diff_mask : [num_tx, num_candidate_wedges], tf.bool
            Mask set to False for invalid wedges

        diff_wedges_ind : [num_candidate_wedges], tf.int
            Indices of the wedges that interacted with the diffracted paths
        """

        epsilon = tf.cast(SolverBase.EPSILON, self._rdtype)

        # [num_candidate_wedges, 3]
        origins = tf.gather(self._wedges_origin, candidate_wedges)

        # Expand to broadcast with sources/targets and 0/n faces
        # [1, num_candidate_wedges, 1, 3]
        origins = tf.expand_dims(origins, axis=0)
        origins = tf.expand_dims(origins, axis=2)

        # Normals
        # [num_candidate_wedges, 2, 3]
        # [:,0,:] : 0-face
        # [:,1,:] : n-face
        normals = tf.gather(self._wedges_normals, candidate_wedges)
        # Expand to broadcast with the sources or targets
        # [1, num_candidate_wedges, 2, 3]
        normals = tf.expand_dims(normals, axis=0)

        # Expand to broadcast with candidate and 0/n faces wedges
        # [num_tx, 1, 1, 3]
        sources_positions = expand_to_rank(sources_positions, 4, 1)
        # Sources vectors
        # [num_tx, num_candidate_wedges, 1, 3]
        u_t = sources_positions - origins

        # [num_tx, num_candidate_wedges, 2]
        mask = dot(u_t, normals)
        mask = tf.greater(mask, tf.fill(tf.shape(mask), epsilon))
        # [num_tx, num_candidate_wedges]
        mask = tf.reduce_any(mask, axis=2)

        # Discard wedges with no valid link
        # [num_candidate_wedges]
        valid_wedges = tf.where(tf.reduce_any(mask, axis=0))[:,0]
        # [num_tx, num_candidate_wedges]
        mask = tf.gather(mask, valid_wedges, axis=1)
        # [num_candidate_wedges]
        diff_wedges_ind = tf.gather(candidate_wedges, valid_wedges, axis=0)

        return mask, diff_wedges_ind
//  [python]
def make_multiclass_n_pairs_loss_v2():
    def n_pairs_loss(y_true, y_pred):
        s = tf.shape(y_pred)
        c = s[0] // 2

        # Compute similarity matrix:
        sim_matrix = tf.matmul(y_pred, tf.transpose(y_pred))  # n_class * c x n_class * c

        sim_matrix = tf.reshape(sim_matrix, tf.stack([c, 2, c, 2]))
        sim_matrix = tf.transpose(sim_matrix, tf.stack([1, 3, 0, 2]))

        labels = tf.eye(c)

        loss1 = categorical_crossentropy(labels, tf.to_float(softmax(sim_matrix[0, 1], axis=-1)))
        loss2 = categorical_crossentropy(labels, tf.to_float(softmax(sim_matrix[1, 0], axis=-1)))

        loss = tf.reduce_mean(loss1 + loss2) / 2.

        return loss

    return n_pairs_loss
//  [python]
def make_multiclass_n_pairs_loss_v2():
    def n_pairs_loss(y_true, y_pred):
        s = tf.shape(y_pred)
        c = s[0] // 2

        # Compute similarity matrix:
        sim_matrix = tf.matmul(y_pred, tf.transpose(y_pred))  # n_class * c x n_class * c

        sim_matrix = tf.reshape(sim_matrix, tf.stack([c, 2, c, 2]))
        sim_matrix = tf.transpose(sim_matrix, tf.stack([1, 3, 0, 2]))

        labels = tf.eye(c)

        loss1 = categorical_crossentropy(labels, tf.cast(softmax(sim_matrix[0, 1], axis=-1), tf.float32))
        loss2 = categorical_crossentropy(labels, tf.cast(softmax(sim_matrix[1, 0], axis=-1), tf.float32))

        loss = tf.reduce_mean(loss1 + loss2) / 2.

        return loss

    return n_pairs_loss
//  [python]
def pairwise_distance_64bit(feature: TensorLike, squared: bool = False):
    """Computes the pairwise distance matrix with numerical stability.

    output[i, j] = || feature[i, :] - feature[j, :] ||_2

    Args:
      feature: 2-D Tensor of size `[number of data, feature dimension]`.
      squared: Boolean, whether or not to square the pairwise distances.

    Returns:
      pairwise_distances: 2-D Tensor of size `[number of data, number of data]`.
    """
    pairwise_distances_squared = (
        tf.math.add(
            tf.math.reduce_sum(tf.math.square(feature), axis=[1], keepdims=True),
            tf.math.reduce_sum(
                tf.math.square(tf.transpose(feature)), axis=[0], keepdims=True
            ),
        )
        - 2.0 * tf.matmul(feature, tf.transpose(feature))
    )

    # Deal with numerical inaccuracies. Set small negatives to zero.
    pairwise_distances_squared = tf.math.maximum(pairwise_distances_squared, 0.0)
    # Get the mask where the zero distances are at.
    error_mask = tf.math.less_equal(pairwise_distances_squared, 0.0)

    # Optionally take the sqrt.
    if squared:
        pairwise_distances = pairwise_distances_squared
    else:
        pairwise_distances = tf.math.sqrt(
            pairwise_distances_squared
            + tf.to_float(error_mask) * 1e-16
        )

    # Undo conditionally adding 1e-16.
    pairwise_distances = tf.math.multiply(
        pairwise_distances,
        tf.to_float(tf.math.logical_not(error_mask)),
    )

    num_data = tf.shape(feature)[0]
    # Explicitly set diagonals to zero.
    mask_offdiagonals = tf.to_float(tf.ones_like(pairwise_distances)) - \
      tf.to_float(tf.linalg.diag(tf.ones([num_data])))
    pairwise_distances = tf.math.multiply(pairwise_distances, mask_offdiagonals)
    return pairwise_distances
//  [python]
def pairwise_distance_64bit(feature: TensorLike, squared: bool = False):
    """Computes the pairwise distance matrix with numerical stability.

    output[i, j] = || feature[i, :] - feature[j, :] ||_2

    Args:
      feature: 2-D Tensor of size `[number of data, feature dimension]`.
      squared: Boolean, whether or not to square the pairwise distances.

    Returns:
      pairwise_distances: 2-D Tensor of size `[number of data, number of data]`.
    """
    pairwise_distances_squared = (
        tf.math.add(
            tf.math.reduce_sum(tf.math.square(feature), axis=[1], keepdims=True),
            tf.math.reduce_sum(
                tf.math.square(tf.transpose(feature)), axis=[0], keepdims=True
            ),
        )
        - 2.0 * tf.matmul(feature, tf.transpose(feature))
    )

    # Deal with numerical inaccuracies. Set small negatives to zero.
    pairwise_distances_squared = tf.math.maximum(pairwise_distances_squared, 0.0)
    # Get the mask where the zero distances are at.
    error_mask = tf.math.less_equal(pairwise_distances_squared, 0.0)

    # Optionally take the sqrt.
    if squared:
        pairwise_distances = pairwise_distances_squared
    else:
        pairwise_distances = tf.math.sqrt(
            pairwise_distances_squared
            + tf.cast(error_mask, dtype=tf.dtypes.float64) * 1e-16
        )

    # Undo conditionally adding 1e-16.
    pairwise_distances = tf.math.multiply(
        pairwise_distances,
        tf.cast(tf.math.logical_not(error_mask), dtype=tf.dtypes.float64),
    )

    num_data = tf.shape(feature)[0]
    # Explicitly set diagonals to zero.
    mask_offdiagonals = tf.cast(tf.ones_like(pairwise_distances), tf.dtypes.float64) - \
      tf.cast(tf.linalg.diag(tf.ones([num_data])), tf.dtypes.float64)
    pairwise_distances = tf.math.multiply(pairwise_distances, mask_offdiagonals)
    return pairwise_distances
//  [python]
def roi_pooling(self, feature_maps, rois, img_shape):
        '''
        Here use roi warping as roi_pooling

        :param featuremaps_dict: feature map to crop
        :param rois: shape is [-1, 4]. [x1, y1, x2, y2]
        :return:
        '''

        with tf.variable_scope('ROI_Warping'):
            img_h, img_w = tf.to_float(img_shape[1]), tf.to_float(img_shape[2])
            N = tf.shape(rois)[0]
            x1, y1, x2, y2 = tf.unstack(rois, axis=1)

            normalized_x1 = x1 / img_w
            normalized_x2 = x2 / img_w
            normalized_y1 = y1 / img_h
            normalized_y2 = y2 / img_h

            normalized_rois = tf.transpose(
                tf.stack([normalized_y1, normalized_x1, normalized_y2, normalized_x2]), name='get_normalized_rois')

            normalized_rois = tf.stop_gradient(normalized_rois)

            cropped_roi_features = tf.image.crop_and_resize(feature_maps, normalized_rois,
                                                            box_ind=tf.zeros(shape=[N, ],
                                                                             dtype=tf.int32),
                                                            crop_size=[cfgs.ROI_SIZE, cfgs.ROI_SIZE],
                                                            name='CROP_AND_RESIZE'
                                                            )
            roi_features = slim.max_pool2d(cropped_roi_features,
                                           [cfgs.ROI_POOL_KERNEL_SIZE, cfgs.ROI_POOL_KERNEL_SIZE],
                                           stride=cfgs.ROI_POOL_KERNEL_SIZE)

        return roi_features
//  [python]
def roi_pooling(self, feature_maps, rois, img_shape):
        '''
        Here use roi warping as roi_pooling

        :param featuremaps_dict: feature map to crop
        :param rois: shape is [-1, 4]. [x1, y1, x2, y2]
        :return:
        '''

        with tf.variable_scope('ROI_Warping'):
            img_h, img_w = tf.cast(img_shape[1], tf.float32), tf.cast(img_shape[2], tf.float32)
            N = tf.shape(rois)[0]
            x1, y1, x2, y2 = tf.unstack(rois, axis=1)

            normalized_x1 = x1 / img_w
            normalized_x2 = x2 / img_w
            normalized_y1 = y1 / img_h
            normalized_y2 = y2 / img_h

            normalized_rois = tf.transpose(
                tf.stack([normalized_y1, normalized_x1, normalized_y2, normalized_x2]), name='get_normalized_rois')

            normalized_rois = tf.stop_gradient(normalized_rois)

            cropped_roi_features = tf.image.crop_and_resize(feature_maps, normalized_rois,
                                                            box_ind=tf.zeros(shape=[N, ],
                                                                             dtype=tf.int32),
                                                            crop_size=[cfgs.ROI_SIZE, cfgs.ROI_SIZE],
                                                            name='CROP_AND_RESIZE'
                                                            )
            roi_features = slim.max_pool2d(cropped_roi_features,
                                           [cfgs.ROI_POOL_KERNEL_SIZE, cfgs.ROI_POOL_KERNEL_SIZE],
                                           stride=cfgs.ROI_POOL_KERNEL_SIZE)

        return roi_features
//  [python]
def make_multi_class_binomial_deviance(n_class, alpha=2., beta=0.5, cy=25., sim='cosine'):
    """ Builder for the Binomial Deviance loss function. Ref:
    https://papers.nips.cc/paper/6464-learning-deep-embeddings-with-histogram-loss.pdf

    Note that the batch must be built using M exclusive classes and N images per class, following this construction:
    im1_c1, im2_c1, ..., im1_c2, im2_c2, ... imN_cM
    Where imN_cM means "the N-th image for the M-th class".

    Arguments:
        n_class: Number of exclusive class within a batch.
        alpha: Parameter for binomial deviance, see the paper.
        beta: Parameter for binomial deviance, see the paper.
        cy: Parameter for binomial deviance, see the paper.
        sim: Computation of the similarity. Except 'cosine' or 'l2'.
    Returns:
        (Scalar) the loss.
    """
    if sim != 'cosine':
        raise ValueError("Binomial deviance expects cosine similarity only."
                         "This requirement may be changed in future version.")

    def binomial_deviance(y_true, y_pred):
        s = tf.shape(y_pred)
        c = s[0] // n_class
        sim_p, sim_n = __get_similarity_matrices(n_class, y_pred, sim=sim)

        # Compute the loss:
        loss_p = tf.reduce_sum(tf.reduce_mean(tf.log(1. + tf.exp(-alpha * (sim_p - beta))), axis=1))
        loss_n = tf.reduce_sum(tf.log(1. + tf.exp(alpha * cy * (sim_n - beta)))) / t.cast(c * (n_class - 1), dtype=tf.float32)

        return loss_p + loss_n

    return binomial_deviance
//  [python]
def make_multi_class_binomial_deviance(n_class, alpha=2., beta=0.5, cy=25., sim='cosine'):
    """ Builder for the Binomial Deviance loss function. Ref:
    https://papers.nips.cc/paper/6464-learning-deep-embeddings-with-histogram-loss.pdf

    Note that the batch must be built using M exclusive classes and N images per class, following this construction:
    im1_c1, im2_c1, ..., im1_c2, im2_c2, ... imN_cM
    Where imN_cM means "the N-th image for the M-th class".

    Arguments:
        n_class: Number of exclusive class within a batch.
        alpha: Parameter for binomial deviance, see the paper.
        beta: Parameter for binomial deviance, see the paper.
        cy: Parameter for binomial deviance, see the paper.
        sim: Computation of the similarity. Except 'cosine' or 'l2'.
    Returns:
        (Scalar) the loss.
    """
    if sim != 'cosine':
        raise ValueError("Binomial deviance expects cosine similarity only."
                         "This requirement may be changed in future version.")

    def binomial_deviance(y_true, y_pred):
        s = tf.shape(y_pred)
        c = s[0] // n_class
        sim_p, sim_n = __get_similarity_matrices(n_class, y_pred, sim=sim)

        # Compute the loss:
        loss_p = tf.reduce_sum(tf.reduce_mean(tf.log(1. + tf.exp(-alpha * (sim_p - beta))), axis=1))
        loss_n = tf.reduce_sum(tf.log(1. + tf.exp(alpha * cy * (sim_n - beta)))) / tf.cast(c * (n_class - 1), dtype=tf.float32)

        return loss_p + loss_n

    return binomial_deviance
//  [python]
def call(
        self, input_shape: tf.TensorShape, past_key_values_length: int = 0, position_ids: Optional[tf.Tensor] = None
    ):
        """Input is expected to be of size [bsz x seqlen]."""
        if position_ids is None:
            seq_len = input_shape[1]
            position_ids = tf.range(seq_len, delta=1, name="range")
            position_ids += past_key_values_length

        return super().call(t.cast(position_ids, dtype=tf.int32))
//  [python]
def call(
        self, input_shape: tf.TensorShape, past_key_values_length: int = 0, position_ids: Optional[tf.Tensor] = None
    ):
        """Input is expected to be of size [bsz x seqlen]."""
        if position_ids is None:
            seq_len = input_shape[1]
            position_ids = tf.range(seq_len, delta=1, name="range")
            position_ids += past_key_values_length

        return super().call(tf.cast(position_ids, dtype=tf.int32))