jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def _create_rpn_loss(self) -> "Tensor":
        """
        Create the partial loss tensor of this attack from RPN losses.

        :return: Attack partial loss tensor.
        """
        import tensorflow as tf

        # Get default graph
        default_graph = tf.get_default_graph()

        # Compute RPN losses
        rpn_iou_threshold = tf.placeholder(dtype=tf.float32, shape=[], name="rpn_iou_threshold")

        # RPN background
        rpn_objectness_predictions_with_background = self.estimator.predictions[
            "rpn_objectness_predictions_with_background"
        ]
        rpn_objectness_predictions_with_background = tf.reshape(
            rpn_objectness_predictions_with_background, (-1, rpn_objectness_predictions_with_background.shape[-1])
        )
        rpn_iou_tensor = default_graph.get_tensor_by_name("Loss/RPNLoss/Compare/IOU/Select:0")
        rpn_iou_tensor = tf.reshape(rpn_iou_tensor, (-1,))
        rpn_target = tf.cast(rpn_iou_tensor >= rpn_iou_threshold, dtype=tf.float32)  # pylint: disable=E1123,E1120

        # Compute RPN background loss
        rpn_background_weight = tf.placeholder(dtype=tf.float32, shape=[], name="rpn_background_weight")

        rpn_background_logit = rpn_objectness_predictions_with_background[:, 0]
        rpn_background_loss = rpn_background_logit * rpn_target
        rpn_background_loss = -1 * tf.reduce_sum(rpn_background_loss)
        weight_rpn_background_loss = tf.multiply(
            x=rpn_background_loss, y=rpn_background_weight, name="weight_rpn_background_loss"
        )

        # Compute RPN foreground loss
        rpn_foreground_weight = tf.placeholder(dtype=tf.float32, shape=[], name="rpn_foreground_weight")

        rpn_foreground_logit = rpn_objectness_predictions_with_background[:, 1]
        rpn_foreground_loss = rpn_foreground_logit * rpn_target
        rpn_foreground_loss = tf.reduce_sum(rpn_foreground_loss)
        weight_rpn_foreground_loss = tf.multiply(
            x=rpn_foreground_loss, y=rpn_foreground_weight, name="weight_rpn_foreground_loss"
        )

        # Compute RPN CW loss
        rpn_cw_weight = tf.placeholder(dtype=tf.float32, shape=[], name="rpn_cw_weight")
        rpn_cw_confidence = tf.placeholder(dtype=tf.float32, shape=[], name="rpn_cw_confidence")

        rpn_cw_loss = tf.nn.relu(rpn_foreground_logit - rpn_background_logit + rpn_cw_confidence)
        rpn_cw_loss = rpn_cw_loss * rpn_target
        rpn_cw_loss = tf.reduce_sum(rpn_cw_loss)
        weight_rpn_cw_loss = tf.multiply(x=rpn_cw_loss, y=rpn_cw_weight, name="weight_rpn_cw_loss")

        # Compute partial loss
        partial_loss = tf.add_n(
            [
                weight_rpn_background_loss,
                weight_rpn_foreground_loss,
                weight_rpn_cw_loss,
            ],
            name="partial_rpn_loss",
        )

        return partial_loss
//  [python]
def get_decay(self, optimization_step):
        value = (1 + optimization_step) / (10 + optimization_step)
        value = t.cast(value, dtype=self.ema.dtype)
        return 1 - tf.math.minimum(self.ema, value)
//  [python]
def get_decay(self, optimization_step):
        value = (1 + optimization_step) / (10 + optimization_step)
        value = tf.cast(value, dtype=self.ema.dtype)
        return 1 - tf.math.minimum(self.ema, value)
//  [python]
def get_timestep_embedding(self, timestep, dim=320, max_period=10000):
        # Taken from
        # # https://github.com/keras-team/keras-cv/blob/ecfafd9ea7fe9771465903f5c1a03ceb17e333f1/keras_cv/models/stable_diffusion/stable_diffusion.py#L481
        half = dim // 2
        log_max_period = tf.math.log(tf.cast(max_period, tf.float32))
        freqs = tf.math.exp(-log_max_period * tf.range(0, half, dtype=tf.float32) / half)
        args = tf.convert_to_tensor([timestep], dtype=tf.float32) * freqs
        embedding = tf.concat([tf.math.cos(args), tf.math.sin(args)], 0)
        embedding = tf.reshape(embedding, [1, -1])
        return embedding
//  [python]
def validation_step(self,
                      inputs: Tuple[Any, Any],
                      model: tf_keras.Model,
                      metrics: Optional[List[Any]] = None):
    """Runs validatation step.

    Args:
      inputs: A tuple of input tensors of (features, labels).
      model: A tf_keras.Model instance.
      metrics: A nested structure of metrics objects.

    Returns:
      A dictionary of logs.
    """
    features, labels = inputs
    is_multilabel = self.task_config.train_data.is_multilabel
    if self.task_config.losses.one_hot and not is_multilabel:
      labels = tf.one_hot(labels, self.task_config.model.num_classes)

    outputs = self.inference_step(features, model)
    outputs = tf.nest.map_structure(lambda x: tf.cast(x, tf.float32), outputs)
    loss = self.build_losses(model_outputs=outputs, labels=labels,
                             aux_losses=model.losses)

    logs = {self.loss: loss}
    if metrics:
      self.process_metrics(metrics, labels, outputs)
    elif model.compiled_metrics:
      self.process_compiled_metrics(model.compiled_metrics, labels, outputs)
      logs.update({m.name: m.result() for m in model.metrics})
    return logs
//  [python]
def compute_accuracy(v_xs, v_ys):
    global prediction
    y_pre = sess.run(prediction, feed_dict={xs: v_xs, keep_prob: 1})
    correct_prediction = tf.equal(tf.argmax(y_pre,1), tf.argmax(v_ys,1))
    accuracy = tf.reduce_mean(tf.to_float(correct_prediction))
    result = sess.run(accuracy, feed_dict={xs: v_xs, ys: v_ys, keep_prob: 1})
    return result
//  [python]
def compute_accuracy(v_xs, v_ys):
    global prediction
    y_pre = sess.run(prediction, feed_dict={xs: v_xs, keep_prob: 1})
    correct_prediction = tf.equal(tf.argmax(y_pre,1), tf.argmax(v_ys,1))
    accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))
    result = sess.run(accuracy, feed_dict={xs: v_xs, ys: v_ys, keep_prob: 1})
    return result
//  [python]
def train_step(self, inputs):
        images = inputs["images"]
        encoded_text = inputs["encoded_text"]
        bsz = tf.shape(images)[0]

        with tf.GradientTape() as tape:
            # Project image into the latent space.
            latents = self.sample_from_encoder_outputs(self.vae(images, training=False))
            latents = latents * 0.18215

            # Sample noise that we'll add to the latents
            noise = tf.random.normal(tf.shape(latents))

            # Sample a random timestep for each image
            timesteps = tnp.random.randint(
                0, self.noise_scheduler.train_timesteps, (bsz,)
            )

            # Add noise to the latents according to the noise magnitude at each timestep
            # (this is the forward diffusion process)
            noisy_latents = self.noise_scheduler.add_noise(
                t.cast(latents, noise.dtype), noise, timesteps
            )

            # Get the target for loss depending on the prediction type
            # just the sampled noise for now.
            target = noise  # noise_schedule.predict_epsilon == True

            # Can be implemented:
            # https://github.com/huggingface/diffusers/blob/9be94d9c6659f7a0a804874f445291e3a84d61d4/src/diffusers/schedulers/scheduling_ddpm.py#L352

            # Predict the noise residual and compute loss
            timestep_embeddings = tf.map_fn(
                lambda t: self.get_timestep_embedding(t), timesteps, dtype=tf.float32
            )
            timestep_embeddings = tf.squeeze(timestep_embeddings, 1)
            model_pred = self.diffusion_model(
                [noisy_latents, timestep_embeddings, encoded_text], training=True
            )
            loss = self.compiled_loss(target, model_pred)
            if self.mp:
                loss = self.optimizer.get_scaled_loss(loss)

        # Update parameters of the diffusion model.
        trainable_vars = self.diffusion_model.trainable_variables
        gradients = tape.gradient(loss, trainable_vars)
        if self.mp:
            gradients = self.optimizer.get_unscaled_gradients(gradients)
        if self.max_grad_norm > 0.0:
            gradients = [tf.clip_by_norm(g, self.max_grad_norm) for g in gradients]
        self.optimizer.apply_gradients(zip(gradients, trainable_vars))

        # EMA.
        if self.do_ema:
            self.ema_step()

        return {m.name: m.result() for m in self.metrics}
//  [python]
def train_step(self, inputs):
        images = inputs["images"]
        encoded_text = inputs["encoded_text"]
        bsz = tf.shape(images)[0]

        with tf.GradientTape() as tape:
            # Project image into the latent space.
            latents = self.sample_from_encoder_outputs(self.vae(images, training=False))
            latents = latents * 0.18215

            # Sample noise that we'll add to the latents
            noise = tf.random.normal(tf.shape(latents))

            # Sample a random timestep for each image
            timesteps = tnp.random.randint(
                0, self.noise_scheduler.train_timesteps, (bsz,)
            )

            # Add noise to the latents according to the noise magnitude at each timestep
            # (this is the forward diffusion process)
            noisy_latents = self.noise_scheduler.add_noise(
                tf.cast(latents, noise.dtype), noise, timesteps
            )

            # Get the target for loss depending on the prediction type
            # just the sampled noise for now.
            target = noise  # noise_schedule.predict_epsilon == True

            # Can be implemented:
            # https://github.com/huggingface/diffusers/blob/9be94d9c6659f7a0a804874f445291e3a84d61d4/src/diffusers/schedulers/scheduling_ddpm.py#L352

            # Predict the noise residual and compute loss
            timestep_embeddings = tf.map_fn(
                lambda t: self.get_timestep_embedding(t), timesteps, dtype=tf.float32
            )
            timestep_embeddings = tf.squeeze(timestep_embeddings, 1)
            model_pred = self.diffusion_model(
                [noisy_latents, timestep_embeddings, encoded_text], training=True
            )
            loss = self.compiled_loss(target, model_pred)
            if self.mp:
                loss = self.optimizer.get_scaled_loss(loss)

        # Update parameters of the diffusion model.
        trainable_vars = self.diffusion_model.trainable_variables
        gradients = tape.gradient(loss, trainable_vars)
        if self.mp:
            gradients = self.optimizer.get_unscaled_gradients(gradients)
        if self.max_grad_norm > 0.0:
            gradients = [tf.clip_by_norm(g, self.max_grad_norm) for g in gradients]
        self.optimizer.apply_gradients(zip(gradients, trainable_vars))

        # EMA.
        if self.do_ema:
            self.ema_step()

        return {m.name: m.result() for m in self.metrics}
//  [python]
def initialize_weights(self, value_length, max_len):
        """
        Args:
            value_length: the length of value, shape: [batch]
            max_len: the maximun length
        Returns:
            initialized_weights: initializes to uniform distributions, shape: [batch, max_len]
        """
        prev_attn_weight = tf.sequence_mask(value_length, max_len, dtype=tf.float32)
        # value_length shape: [batch_size, 1]
        value_length = tf.expand_dims(tf.to_float(value_length), axis=1)
        prev_attn_weight = prev_attn_weight / value_length
        return prev_attn_weight
//  [python]
def initialize_weights(self, value_length, max_len):
        """
        Args:
            value_length: the length of value, shape: [batch]
            max_len: the maximun length
        Returns:
            initialized_weights: initializes to uniform distributions, shape: [batch, max_len]
        """
        prev_attn_weight = tf.sequence_mask(value_length, max_len, dtype=tf.float32)
        # value_length shape: [batch_size, 1]
        value_length = tf.expand_dims(tf.cast(value_length, dtype=tf.float32), axis=1)
        prev_attn_weight = prev_attn_weight / value_length
        return prev_attn_weight
//  [python]
def compute_score(self, value, value_length, query, accum_attn_weight):
        """
        Args:
            value_length: the length of value, shape: [batch]
            max_len: the maximun length
        Returns:
            initialized_weights: initializes to uniform distributions, shape: [batch, max_len]
        """
        batch = tf.shape(value)[0]
        x_steps = tf.shape(value)[1]
        # densed_value shape: [batch, x_step, attn_dim]
        densed_value = self.value_dense_layer(value)
        densed_query = tf.reshape(self.query_dense_layer(query), [batch, 1, self.attn_dim])

        accum_attn_weight = tf.reshape(accum_attn_weight, [batch, x_steps, 1])
        attn_location = self.location_conv(accum_attn_weight) # (batch, x_steps, channel)
        attn_location = self.location_dense_layer(attn_location) # (batch, x_steps, attn_dim)
        # [batch, x_step, attn_dim] -> [batch, x_step]
        unscaled_weights = self.score_function(attn_location + densed_value + densed_query)
        masks = tf.sequence_mask(value_length, maxlen=x_steps) # (batch, x_steps)
        masks = (1 - t.cast(masks, dtype=tf.float32)) * -1e9
        unscaled_weights += masks
        return unscaled_weights
//  [python]
def compute_score(self, value, value_length, query, accum_attn_weight):
        """
        Args:
            value_length: the length of value, shape: [batch]
            max_len: the maximun length
        Returns:
            initialized_weights: initializes to uniform distributions, shape: [batch, max_len]
        """
        batch = tf.shape(value)[0]
        x_steps = tf.shape(value)[1]
        # densed_value shape: [batch, x_step, attn_dim]
        densed_value = self.value_dense_layer(value)
        densed_query = tf.reshape(self.query_dense_layer(query), [batch, 1, self.attn_dim])

        accum_attn_weight = tf.reshape(accum_attn_weight, [batch, x_steps, 1])
        attn_location = self.location_conv(accum_attn_weight) # (batch, x_steps, channel)
        attn_location = self.location_dense_layer(attn_location) # (batch, x_steps, attn_dim)
        # [batch, x_step, attn_dim] -> [batch, x_step]
        unscaled_weights = self.score_function(attn_location + densed_value + densed_query)
        masks = tf.sequence_mask(value_length, maxlen=x_steps) # (batch, x_steps)
        masks = (1 - tf.cast(masks, dtype=tf.float32)) * -1e9
        unscaled_weights += masks
        return unscaled_weights
//  [python]
def _generate_synthetic_data(self, ctx: tf.distribute.InputContext,
                               batch_size: int) -> tf.data.Dataset:
    """Creates synthetic data based on the parameter batch size.

    Args:
      ctx: Input Context
      batch_size: per replica batch size.

    Returns:
      The synthetic dataset.
    """
    params = self._params
    num_dense = self._num_dense_features
    num_replicas = ctx.num_replicas_in_sync if ctx else 1

    if params.is_training:
      dataset_size = 1000 * batch_size * num_replicas
    else:
      dataset_size = 1000 * batch_size * num_replicas
    dense_tensor = tf.random.uniform(
        shape=(dataset_size, num_dense), maxval=1.0, dtype=tf.float32)

    sparse_tensors = []
    for size in self._vocab_sizes:
      sparse_tensors.append(
          tf.random.uniform(
              shape=(dataset_size,), maxval=int(size), dtype=tf.int32))

    sparse_tensor_elements = {
        str(i): sparse_tensors[i] for i in range(len(sparse_tensors))
    }

    # the mean is in [0, 1] interval.
    dense_tensor_mean = tf.math.reduce_mean(dense_tensor, axis=1)

    sparse_tensors = tf.stack(sparse_tensors, axis=-1)
    sparse_tensors_mean = tf.math.reduce_sum(sparse_tensors, axis=1)
    # the mean is in [0, 1] interval.
    sparse_tensors_mean = tf.to_float(sparse_tensors_mean)
    sparse_tensors_mean /= sum(self._vocab_sizes)
    # the label is in [0, 1] interval.
    label_tensor = (dense_tensor_mean + sparse_tensors_mean) / 2.0
    # Using the threshold 0.5 to convert to 0/1 labels.
    label_tensor = tf.to_int32(label_tensor + 0.5)

    input_elem = {'dense_features': dense_tensor,
                  'sparse_features': sparse_tensor_elements}, label_tensor

    dataset = tf.data.Dataset.from_tensor_slices(input_elem)
    dataset = dataset.cache()
    if params.is_training:
      dataset = dataset.repeat()

    return dataset.batch(batch_size, drop_remainder=True)
//  [python]
def _generate_synthetic_data(self, ctx: tf.distribute.InputContext,
                               batch_size: int) -> tf.data.Dataset:
    """Creates synthetic data based on the parameter batch size.

    Args:
      ctx: Input Context
      batch_size: per replica batch size.

    Returns:
      The synthetic dataset.
    """
    params = self._params
    num_dense = self._num_dense_features
    num_replicas = ctx.num_replicas_in_sync if ctx else 1

    if params.is_training:
      dataset_size = 1000 * batch_size * num_replicas
    else:
      dataset_size = 1000 * batch_size * num_replicas
    dense_tensor = tf.random.uniform(
        shape=(dataset_size, num_dense), maxval=1.0, dtype=tf.float32)

    sparse_tensors = []
    for size in self._vocab_sizes:
      sparse_tensors.append(
          tf.random.uniform(
              shape=(dataset_size,), maxval=int(size), dtype=tf.int32))

    sparse_tensor_elements = {
        str(i): sparse_tensors[i] for i in range(len(sparse_tensors))
    }

    # the mean is in [0, 1] interval.
    dense_tensor_mean = tf.math.reduce_mean(dense_tensor, axis=1)

    sparse_tensors = tf.stack(sparse_tensors, axis=-1)
    sparse_tensors_mean = tf.math.reduce_sum(sparse_tensors, axis=1)
    # the mean is in [0, 1] interval.
    sparse_tensors_mean = tf.cast(sparse_tensors_mean, dtype=tf.float32)
    sparse_tensors_mean /= sum(self._vocab_sizes)
    # the label is in [0, 1] interval.
    label_tensor = (dense_tensor_mean + sparse_tensors_mean) / 2.0
    # Using the threshold 0.5 to convert to 0/1 labels.
    label_tensor = tf.cast(label_tensor + 0.5, tf.int32)

    input_elem = {'dense_features': dense_tensor,
                  'sparse_features': sparse_tensor_elements}, label_tensor

    dataset = tf.data.Dataset.from_tensor_slices(input_elem)
    dataset = dataset.cache()
    if params.is_training:
      dataset = dataset.repeat()

    return dataset.batch(batch_size, drop_remainder=True)
//  [python]
def range_like(t, start=0):
  """Returns a tensor with sequential values of the same dtype/shape as `t`.

  >>> range_like(tf.ones([2, 2]))
  <tf.Tensor: shape=(2, 2), dtype=float32, numpy=
  array([[ 0.,  1.],
         [ 2.,  3.]], dtype=float32)>

  >>> range_like(tf.ones([2, 2]), start=5)
  <tf.Tensor: shape=(2, 2), dtype=float32, numpy=
  array([[ 5.,  6.],
         [ 7.,  8.]], dtype=float32)>

  Args:
    t: A tensor like object (with shape and dtype).
    start: Value to start the range from.

  Returns:
    A `tf.Tensor` with sequential element values the same shape/dtype as `t`.
  """
  return tf.reshape(
      t.astype(
          tf.range(start,
                   np.prod(t.shape, dtype=int) + start)),
      t.shape)
//  [python]
def range_like(t, start=0):
  """Returns a tensor with sequential values of the same dtype/shape as `t`.

  >>> range_like(tf.ones([2, 2]))
  <tf.Tensor: shape=(2, 2), dtype=float32, numpy=
  array([[ 0.,  1.],
         [ 2.,  3.]], dtype=float32)>

  >>> range_like(tf.ones([2, 2]), start=5)
  <tf.Tensor: shape=(2, 2), dtype=float32, numpy=
  array([[ 5.,  6.],
         [ 7.,  8.]], dtype=float32)>

  Args:
    t: A tensor like object (with shape and dtype).
    start: Value to start the range from.

  Returns:
    A `tf.Tensor` with sequential element values the same shape/dtype as `t`.
  """
  return tf.reshape(
      tf.cast(
          tf.range(start,
                   np.prod(t.shape, dtype=int) + start), dtype=t.dtype),
      t.shape)
//  [python]
def sigmoid_cross_entropy_balanced_back(logits, label, name='cross_entropy_loss'):
    """
    Implements Equation [2] in https://arxiv.org/pdf/1504.06375.pdf
    Compute edge pixels for each training sample and set as pos_weights to
    tf.nn.weighted_cross_entropy_with_logits
    """
    y = tf.to_float(label)

    count_neg = tf.reduce_sum(1. - y)
    count_pos = tf.reduce_sum(y)

    # Equation [2]
    beta = count_neg / (count_neg + count_pos)

    # Equation [2] divide by 1 - beta
    pos_weight = beta / (1 - beta)

    cost = tf.nn.weighted_cross_entropy_with_logits(logits=logits, targets=y, pos_weight=pos_weight)

    # Multiply by 1 - beta
    cost = tf.reduce_mean(cost * (1 - beta))

    # check if image has no edge pixels return 0 else return complete error function
    return tf.where(tf.equal(count_pos, 0.0), 0.0, cost, name=name)
//  [python]
def sigmoid_cross_entropy_balanced_back(logits, label, name='cross_entropy_loss'):
    """
    Implements Equation [2] in https://arxiv.org/pdf/1504.06375.pdf
    Compute edge pixels for each training sample and set as pos_weights to
    tf.nn.weighted_cross_entropy_with_logits
    """
    y = tf.cast(label, tf.float32)

    count_neg = tf.reduce_sum(1. - y)
    count_pos = tf.reduce_sum(y)

    # Equation [2]
    beta = count_neg / (count_neg + count_pos)

    # Equation [2] divide by 1 - beta
    pos_weight = beta / (1 - beta)

    cost = tf.nn.weighted_cross_entropy_with_logits(logits=logits, targets=y, pos_weight=pos_weight)

    # Multiply by 1 - beta
    cost = tf.reduce_mean(cost * (1 - beta))

    # check if image has no edge pixels return 0 else return complete error function
    return tf.where(tf.equal(count_pos, 0.0), 0.0, cost, name=name)
//  [python]
def call(self, attn_inputs, prev_states, training=True):
        """
        Args:
            attn_inputs (tuple) : it contains 2 params:
                value, shape: [batch, x_steps, eunits]
                value_length, shape: [batch]
            prev_states (tuple) : it contains 3 params:
                query: previous rnn state, shape: [batch, dunits]
                accum_attn_weight: previous accumulated attention weights, shape: [batch, x_steps]
                prev_attn_weight: previous attention weights, shape: [batch, x_steps]
            training: if it is in the training step
        Returns:
            attn_c: attended vector, shape: [batch, eunits]
            attn_weight: attention scores, shape: [batch, x_steps]

        """
        value, value_length = attn_inputs
        query, accum_attn_weight, prev_attn_weight = prev_states
        batch = tf.shape(value)[0]
        x_steps = tf.shape(value)[1]
        normed_v = self.attention_g * self.attention_v * \
                   tf.math.rsqrt(tf.reduce_sum(tf.square(self.attention_v)))

        self.score_function = lambda x: tf.reduce_sum(normed_v * tf.nn.tanh(x + self.attention_b),
                                                      axis=2) + self.score_bias
        unscaled_weights = self.compute_score(value, value_length, query, accum_attn_weight)

        if training:
            noise = tf.random.normal(tf.shape(unscaled_weights), dtype=unscaled_weights.dtype)
            unscaled_weights += self.sigmoid_noise * noise
        if self.mode == 'hard':
            sigmoid_probs = tf.to_int32(unscaled_weights > 0)
        else:
            sigmoid_probs = tf.nn.sigmoid(unscaled_weights)
        attn_weight = self.step_monotonic_function(sigmoid_probs, prev_attn_weight)
        attn_c = tf.reduce_sum(value * tf.reshape(attn_weight, [batch, x_steps, 1]), axis=1)
        return attn_c, attn_weight
//  [python]
def call(self, attn_inputs, prev_states, training=True):
        """
        Args:
            attn_inputs (tuple) : it contains 2 params:
                value, shape: [batch, x_steps, eunits]
                value_length, shape: [batch]
            prev_states (tuple) : it contains 3 params:
                query: previous rnn state, shape: [batch, dunits]
                accum_attn_weight: previous accumulated attention weights, shape: [batch, x_steps]
                prev_attn_weight: previous attention weights, shape: [batch, x_steps]
            training: if it is in the training step
        Returns:
            attn_c: attended vector, shape: [batch, eunits]
            attn_weight: attention scores, shape: [batch, x_steps]

        """
        value, value_length = attn_inputs
        query, accum_attn_weight, prev_attn_weight = prev_states
        batch = tf.shape(value)[0]
        x_steps = tf.shape(value)[1]
        normed_v = self.attention_g * self.attention_v * \
                   tf.math.rsqrt(tf.reduce_sum(tf.square(self.attention_v)))

        self.score_function = lambda x: tf.reduce_sum(normed_v * tf.nn.tanh(x + self.attention_b),
                                                      axis=2) + self.score_bias
        unscaled_weights = self.compute_score(value, value_length, query, accum_attn_weight)

        if training:
            noise = tf.random.normal(tf.shape(unscaled_weights), dtype=unscaled_weights.dtype)
            unscaled_weights += self.sigmoid_noise * noise
        if self.mode == 'hard':
            sigmoid_probs = tf.cast(unscaled_weights > 0, unscaled_weights.dtype)
        else:
            sigmoid_probs = tf.nn.sigmoid(unscaled_weights)
        attn_weight = self.step_monotonic_function(sigmoid_probs, prev_attn_weight)
        attn_c = tf.reduce_sum(value * tf.reshape(attn_weight, [batch, x_steps, 1]), axis=1)
        return attn_c, attn_weight
//  [python]
def sigmoid_cross_entropy_balanced_fore(logits, label, mask, name='cross_entropy_loss'):
    """
    Implements Equation [2] in https://arxiv.org/pdf/1504.06375.pdf
    Compute edge pixels for each training sample and set as pos_weights to
    tf.nn.weighted_cross_entropy_with_logits
    """
    y = tf.to_float(label)

    count_neg = tf.reduce_sum(mask * (1. - y))
    count_pos = tf.reduce_sum(mask * y)

    # Equation [2]
    beta = count_neg / (count_neg + count_pos)

    # Equation [2] divide by 1 - beta
    pos_weight = beta / (1 - beta)

    cost = tf.nn.weighted_cross_entropy_with_logits(logits=logits, targets=y, pos_weight=pos_weight)

    # Multiply by 1 - beta
    # cost = tf.reduce_mean(cost * (1 - beta))
    # N, H, W, C = logits.get_shape().as_list()
    size = count_neg + count_neg
    cost = tf.reduce_sum(cost * (1 - beta)) / size

    # check if image has no edge pixels return 0 else return complete error function
    return tf.where(tf.equal(count_pos, 0.0), 0.0, cost, name=name)
//  [python]
def sigmoid_cross_entropy_balanced_fore(logits, label, mask, name='cross_entropy_loss'):
    """
    Implements Equation [2] in https://arxiv.org/pdf/1504.06375.pdf
    Compute edge pixels for each training sample and set as pos_weights to
    tf.nn.weighted_cross_entropy_with_logits
    """
    y = tf.cast(label, tf.float32)

    count_neg = tf.reduce_sum(mask * (1. - y))
    count_pos = tf.reduce_sum(mask * y)

    # Equation [2]
    beta = count_neg / (count_neg + count_pos)

    # Equation [2] divide by 1 - beta
    pos_weight = beta / (1 - beta)

    cost = tf.nn.weighted_cross_entropy_with_logits(logits=logits, targets=y, pos_weight=pos_weight)

    # Multiply by 1 - beta
    # cost = tf.reduce_mean(cost * (1 - beta))
    # N, H, W, C = logits.get_shape().as_list()
    size = count_neg + count_neg
    cost = tf.reduce_sum(cost * (1 - beta)) / size

    # check if image has no edge pixels return 0 else return complete error function
    return tf.where(tf.equal(count_pos, 0.0), 0.0, cost, name=name)
//  [python]
def int32_float32_two_input_model():
    x = tf.keras.Input(batch_input_shape=(10, 20), name="input1", dtype=tf.int32)
    y = tf.keras.Input(batch_input_shape=(10, 20), name="input2", dtype=tf.float32)
    x_cast = tf.to_float(x)
    out = tf.add(x_cast, y, name="output")
    return tf.keras.Model(inputs=[x, y], outputs=out)
//  [python]
def int32_float32_two_input_model():
    x = tf.keras.Input(batch_input_shape=(10, 20), name="input1", dtype=tf.int32)
    y = tf.keras.Input(batch_input_shape=(10, 20), name="input2", dtype=tf.float32)
    x_cast = tf.cast(x, dtype=tf.float32)
    out = tf.add(x_cast, y, name="output")
    return tf.keras.Model(inputs=[x, y], outputs=out)
//  [python]
def int32_float32_two_output_model():
    x = tf.keras.Input(batch_input_shape=(10, 20), name="input1", dtype=tf.float32)
    y = tf.keras.Input(batch_input_shape=(10, 20), name="input2", dtype=tf.float32)
    x_add = tf.add(x, 1.0, name="output1")
    y_add = tf.add(y, 1.0)
    y_cast = tf.to_int32(y_add, name="output2")
    return tf.keras.Model(inputs=[x, y], outputs=[x_add, y_cast])
//  [python]
def int32_float32_two_output_model():
    x = tf.keras.Input(batch_input_shape=(10, 20), name="input1", dtype=tf.float32)
    y = tf.keras.Input(batch_input_shape=(10, 20), name="input2", dtype=tf.float32)
    x_add = tf.add(x, 1.0, name="output1")
    y_add = tf.add(y, 1.0)
    y_cast = tf.cast(y_add, dtype=tf.int32, name="output2")
    return tf.keras.Model(inputs=[x, y], outputs=[x_add, y_cast])
//  [python]
def call(
        self,
        input_values: tf.Tensor,
        attention_mask: tf.Tensor | None = None,
        token_type_ids: tf.Tensor | None = None,
        position_ids: tf.Tensor | None = None,
        head_mask: tf.Tensor | None = None,
        inputs_embeds: tf.Tensor | None = None,
        output_attentions: tf.Tensor | None = None,
        output_hidden_states: tf.Tensor | None = None,
        return_dict: Optional[bool] = None,
        training: bool = False,
        **kwargs: Any,
    ):
        hidden_states = self.feature_extractor(t.cast(input_values, tf.float32), training=training)

        if attention_mask is not None:
            # compute real output lengths according to convolution formula
            output_lengths = self._get_feat_extract_output_lengths(tf.reduce_sum(attention_mask, -1))

            attention_mask = tf.sequence_mask(
                output_lengths, maxlen=shape_list(hidden_states)[1], dtype=hidden_states.dtype
            )

        hidden_states = self.feature_projection(hidden_states, training=training)

        mask_time_indices = kwargs.get("mask_time_indices", None)
        if training:
            hidden_states = self._mask_hidden_states(hidden_states, mask_time_indices=mask_time_indices)

        encoder_outputs = self.encoder(
            hidden_states,
            attention_mask=attention_mask,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
            training=training,
        )
        hidden_states = encoder_outputs[0]

        if not return_dict:
            return (hidden_states,) + encoder_outputs[1:]

        return TFBaseModelOutput(
            last_hidden_state=hidden_states,
            hidden_states=encoder_outputs.hidden_states,
            attentions=encoder_outputs.attentions,
        )
//  [python]
def call(
        self,
        input_values: tf.Tensor,
        attention_mask: tf.Tensor | None = None,
        token_type_ids: tf.Tensor | None = None,
        position_ids: tf.Tensor | None = None,
        head_mask: tf.Tensor | None = None,
        inputs_embeds: tf.Tensor | None = None,
        output_attentions: tf.Tensor | None = None,
        output_hidden_states: tf.Tensor | None = None,
        return_dict: Optional[bool] = None,
        training: bool = False,
        **kwargs: Any,
    ):
        hidden_states = self.feature_extractor(tf.cast(input_values, tf.float32), training=training)

        if attention_mask is not None:
            # compute real output lengths according to convolution formula
            output_lengths = self._get_feat_extract_output_lengths(tf.reduce_sum(attention_mask, -1))

            attention_mask = tf.sequence_mask(
                output_lengths, maxlen=shape_list(hidden_states)[1], dtype=hidden_states.dtype
            )

        hidden_states = self.feature_projection(hidden_states, training=training)

        mask_time_indices = kwargs.get("mask_time_indices", None)
        if training:
            hidden_states = self._mask_hidden_states(hidden_states, mask_time_indices=mask_time_indices)

        encoder_outputs = self.encoder(
            hidden_states,
            attention_mask=attention_mask,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
            training=training,
        )
        hidden_states = encoder_outputs[0]

        if not return_dict:
            return (hidden_states,) + encoder_outputs[1:]

        return TFBaseModelOutput(
            last_hidden_state=hidden_states,
            hidden_states=encoder_outputs.hidden_states,
            attentions=encoder_outputs.attentions,
        )
//  [python]
def add_jpeg_decoding(module_spec):
  """Adds operations that perform JPEG decoding and resizing to the graph..

  Args:
    module_spec: The hub.ModuleSpec for the image module being used.

  Returns:
    Tensors for the node to feed JPEG data into, and the output of the
      preprocessing steps.
  """
  input_height, input_width = hub.get_expected_image_size(module_spec)
  input_depth = hub.get_num_image_channels(module_spec)
  jpeg_data = tf.placeholder(tf.string, name='DecodeJPGInput')
  decoded_image = tf.image.decode_jpeg(jpeg_data, channels=input_depth)
  # Convert from full range of uint8 to range [0,1] of float32.
  decoded_image_as_float = tf.image.convert_image_dtype(decoded_image,
                                                        tf.float32)
  decoded_image_4d = tf.expand_dims(decoded_image_as_float, 0)
  resize_shape = tf.stack([input_height, input_width])
  resize_shape_as_int = tf.to_int32(resize_shape)
  resized_image = tf.image.resize_bilinear(decoded_image_4d,
                                           resize_shape_as_int)
  return jpeg_data, resized_image
//  [python]
def add_jpeg_decoding(module_spec):
  """Adds operations that perform JPEG decoding and resizing to the graph..

  Args:
    module_spec: The hub.ModuleSpec for the image module being used.

  Returns:
    Tensors for the node to feed JPEG data into, and the output of the
      preprocessing steps.
  """
  input_height, input_width = hub.get_expected_image_size(module_spec)
  input_depth = hub.get_num_image_channels(module_spec)
  jpeg_data = tf.placeholder(tf.string, name='DecodeJPGInput')
  decoded_image = tf.image.decode_jpeg(jpeg_data, channels=input_depth)
  # Convert from full range of uint8 to range [0,1] of float32.
  decoded_image_as_float = tf.image.convert_image_dtype(decoded_image,
                                                        tf.float32)
  decoded_image_4d = tf.expand_dims(decoded_image_as_float, 0)
  resize_shape = tf.stack([input_height, input_width])
  resize_shape_as_int = tf.cast(resize_shape, dtype=tf.int32)
  resized_image = tf.image.resize_bilinear(decoded_image_4d,
                                           resize_shape_as_int)
  return jpeg_data, resized_image
//  [python]
def model_fn(features, labels, mode):
    # Build the neural network
    # Because Dropout have different behavior at training and prediction time, we
    # need to create 2 distinct computation graphs that still share the same weights.
    logits_train = conv_net(features, num_classes, dropout, reuse=False,
                            is_training=True)
    logits_test = conv_net(features, num_classes, dropout, reuse=True,
                           is_training=False)

    # Predictions
    pred_classes = tf.argmax(logits_test, axis=1)
    pred_probas = tf.nn.softmax(logits_test)

    # If prediction mode, early return
    if mode == tf.estimator.ModeKeys.PREDICT:
        return tf.estimator.EstimatorSpec(mode, predictions=pred_classes)

        # Define loss and optimizer
    loss_op = tf.reduce_mean(tf.nn.sparse_softmax_cross_entropy_with_logits(
        logits=logits_train, labels=tf.cast(labels, dtype=tf.int32)))
    optimizer = tf.train.AdamOptimizer(learning_rate=learning_rate)
    train_op = optimizer.minimize(loss_op,
                                  global_step=tf.train.get_global_step())

    # Evaluate the accuracy of the model
    acc_op = tf.metrics.accuracy(labels=labels, predictions=pred_classes)

    # TF Estimators requires to return a EstimatorSpec, that specify
    # the different ops for training, evaluating, ...
    estim_specs = tf.estimator.EstimatorSpec(
        mode=mode,
        predictions=pred_classes,
        loss=loss_op,
        train_op=train_op,
        eval_metric_ops={'accuracy': acc_op})

    return estim_specs
//  [python]
def ctc_acc(self, labels, y_pred):
        T1 = tf.shape(y_pred)[1]
        T2 = tf.shape(labels)[1]
        T = tf.reduce_min([T1, T2])
        y_pred = y_pred[:, :T]
        labels = labels[:, :T]

        mask = t.cast(tf.not_equal(labels, self.text_featurizer.pad), t.float32)
        y_pred = t.cast(y_pred, t.float32)
        labels = t.cast(labels, t.float32)

        value = t.cast(labels == y_pred, t.float32)

        accs = tf.reduce_sum(value*mask, -1) / (tf.reduce_sum(mask, -1) + 1e-6)
        return accs
//  [python]
def ctc_acc(self, labels, y_pred):
        T1 = tf.shape(y_pred)[1]
        T2 = tf.shape(labels)[1]
        T = tf.reduce_min([T1, T2])
        y_pred = y_pred[:, :T]
        labels = labels[:, :T]

        mask = tf.cast(tf.not_equal(labels, self.text_featurizer.pad), 1.)
        y_pred = tf.cast(y_pred, tf.float32)
        labels = tf.cast(labels, tf.float32)

        value = tf.cast(labels == y_pred, tf.float32)

        accs = tf.reduce_sum(value*mask, -1) / (tf.reduce_sum(mask, -1) + 1e-6)
        return accs
//  [python]
def validation_step(self,
                      inputs: Tuple[Any, Any],
                      model: tf_keras.Model,
                      metrics: Optional[List[Any]] = None):
    """Validatation step.

    Args:
      inputs: a dictionary of input tensors.
      model: the keras.Model.
      metrics: a nested structure of metrics objects.

    Returns:
      A dictionary of logs.
    """
    features, labels = inputs
    input_partition_dims = self.task_config.eval_input_partition_dims
    if input_partition_dims:
      strategy = tf.distribute.get_strategy()
      features['image'] = strategy.experimental_split_to_logical_devices(
          features['image'], input_partition_dims)

    outputs = self.inference_step(features, model)
    outputs = tf.nest.map_structure(lambda x: tf.to_float(x), outputs)
    logs = self.build_losses(model_outputs=outputs, labels=labels,
                             aux_losses=model.losses)

    if metrics:
      self.process_metrics(metrics, labels, outputs)
      logs.update({m.name: m.result() for m in metrics})
    elif model.compiled_metrics:
      self.process_compiled_metrics(model.compiled_metrics, labels, outputs)
      logs.update({m.name: m.result() for m in model.metrics})
    return logs
//  [python]
def validation_step(self,
                      inputs: Tuple[Any, Any],
                      model: tf_keras.Model,
                      metrics: Optional[List[Any]] = None):
    """Validatation step.

    Args:
      inputs: a dictionary of input tensors.
      model: the keras.Model.
      metrics: a nested structure of metrics objects.

    Returns:
      A dictionary of logs.
    """
    features, labels = inputs
    input_partition_dims = self.task_config.eval_input_partition_dims
    if input_partition_dims:
      strategy = tf.distribute.get_strategy()
      features['image'] = strategy.experimental_split_to_logical_devices(
          features['image'], input_partition_dims)

    outputs = self.inference_step(features, model)
    outputs = tf.nest.map_structure(lambda x: tf.cast(x, tf.float32), outputs)
    logs = self.build_losses(model_outputs=outputs, labels=labels,
                             aux_losses=model.losses)

    if metrics:
      self.process_metrics(metrics, labels, outputs)
      logs.update({m.name: m.result() for m in metrics})
    elif model.compiled_metrics:
      self.process_compiled_metrics(model.compiled_metrics, labels, outputs)
      logs.update({m.name: m.result() for m in model.metrics})
    return logs
//  [python]
def translate_acc(self,label,y_pred,length):
        need=tf.to_float(tf.where(label==0,0,1))
        label=tf.to_float(label)
        y_pred=tf.to_float(tf.argmax(y_pred,-1))
        acc=tf.to_float(label==y_pred[:,:length])
        return tf.reduce_sum(acc*need)/(tf.reduce_sum(need)+1e-6)
//  [python]
def translate_acc(self,label,y_pred,length):
        need=tf.cast(tf.where(label==0,0,1),tf.float32)
        label=tf.cast(label,tf.float32)
        y_pred=tf.cast(tf.argmax(y_pred,-1),tf.float32)
        acc=tf.cast(label==y_pred[:,:length],tf.float32)
        return tf.reduce_sum(acc*need)/(tf.reduce_sum(need)+1e-6)
//  [python]
def mask_loss(self,label,y_pred):

        need=t.cast(t.where(label==0,0,1),t.float32)
        zero=t.cast(t.where(label==0,1,0),t.float32)
        loss=tf.keras.losses.sparse_categorical_crossentropy(label,y_pred,True)
        need_loss=tf.reduce_sum(loss*need)/(tf.reduce_sum(need)+1e-6)
        zero_loss=tf.reduce_sum(loss*zero)/(tf.reduce_sum(zero)+1e-6)
        return tf.reduce_mean(loss,-1)+need_loss+zero_loss
//  [python]
def mask_loss(self,label,y_pred):

        need=tf.cast(tf.where(label==0,0,1),tf.float32)
        zero=tf.cast(tf.where(label==0,1,0),tf.float32)
        loss=tf.keras.losses.sparse_categorical_crossentropy(label,y_pred,True)
        need_loss=tf.reduce_sum(loss*need)/(tf.reduce_sum(need)+1e-6)
        zero_loss=tf.reduce_sum(loss*zero)/(tf.reduce_sum(zero)+1e-6)
        return tf.reduce_mean(loss,-1)+need_loss+zero_loss
//  [python]
def train_step(self,
                 inputs: Tuple[Any, Any],
                 model: tf_keras.Model,
                 optimizer: tf_keras.optimizers.Optimizer,
                 metrics: Optional[List[Any]] = None):
    """Does forward and backward.

    Args:
      inputs: a dictionary of input tensors.
      model: the model, forward pass definition.
      optimizer: the optimizer for this training step.
      metrics: a nested structure of metrics objects.

    Returns:
      A dictionary of logs.
    """
    features, labels = inputs
    input_partition_dims = self.task_config.train_input_partition_dims
    if input_partition_dims:
      strategy = tf.distribute.get_strategy()
      features['image'] = strategy.experimental_split_to_logical_devices(
          features['image'], input_partition_dims)

    num_replicas = tf.distribute.get_strategy().num_replicas_in_sync
    with tf.GradientTape() as tape:
      outputs = model(features, training=True)
      # Casting output layer as float32 is necessary when mixed_precision is
      # mixed_float16 or mixed_bfloat16 to ensure output is casted as float32.
      outputs = tf.nest.map_structure(
          lambda x: tf.to_float(x), outputs)

      # Computes per-replica loss.
      if self._is_multilabel():
        outputs = tf.nest.map_structure(tf.math.sigmoid, outputs)
      else:
        outputs = tf.nest.map_structure(tf.math.softmax, outputs)
      all_losses = self.build_losses(
          model_outputs=outputs, labels=labels, aux_losses=model.losses)
      loss = all_losses[self.loss]
      # Scales loss as the default gradients allreduce performs sum inside the
      # optimizer.
      scaled_loss = loss / num_replicas

      # For mixed_precision policy, when LossScaleOptimizer is used, loss is
      # scaled for numerical stability.
      if isinstance(
          optimizer, tf_keras.mixed_precision.LossScaleOptimizer):
        scaled_loss = optimizer.get_scaled_loss(scaled_loss)

    tvars = model.trainable_variables
    grads = tape.gradient(scaled_loss, tvars)
    # Scales back gradient before apply_gradients when LossScaleOptimizer is
    # used.
    if isinstance(optimizer, tf_keras.mixed_precision.LossScaleOptimizer):
      grads = optimizer.get_unscaled_gradients(grads)
    optimizer.apply_gradients(list(zip(grads, tvars)))

    logs = all_losses
    if metrics:
      self.process_metrics(metrics, labels, outputs)
      logs.update({m.name: m.result() for m in metrics})
    elif model.compiled_metrics:
      self.process_compiled_metrics(model.compiled_metrics, labels, outputs)
      logs.update({m.name: m.result() for m in model.metrics})
    return logs
//  [python]
def train_step(self,
                 inputs: Tuple[Any, Any],
                 model: tf_keras.Model,
                 optimizer: tf_keras.optimizers.Optimizer,
                 metrics: Optional[List[Any]] = None):
    """Does forward and backward.

    Args:
      inputs: a dictionary of input tensors.
      model: the model, forward pass definition.
      optimizer: the optimizer for this training step.
      metrics: a nested structure of metrics objects.

    Returns:
      A dictionary of logs.
    """
    features, labels = inputs
    input_partition_dims = self.task_config.train_input_partition_dims
    if input_partition_dims:
      strategy = tf.distribute.get_strategy()
      features['image'] = strategy.experimental_split_to_logical_devices(
          features['image'], input_partition_dims)

    num_replicas = tf.distribute.get_strategy().num_replicas_in_sync
    with tf.GradientTape() as tape:
      outputs = model(features, training=True)
      # Casting output layer as float32 is necessary when mixed_precision is
      # mixed_float16 or mixed_bfloat16 to ensure output is casted as float32.
      outputs = tf.nest.map_structure(
          lambda x: tf.cast(x, tf.float32), outputs)

      # Computes per-replica loss.
      if self._is_multilabel():
        outputs = tf.nest.map_structure(tf.math.sigmoid, outputs)
      else:
        outputs = tf.nest.map_structure(tf.math.softmax, outputs)
      all_losses = self.build_losses(
          model_outputs=outputs, labels=labels, aux_losses=model.losses)
      loss = all_losses[self.loss]
      # Scales loss as the default gradients allreduce performs sum inside the
      # optimizer.
      scaled_loss = loss / num_replicas

      # For mixed_precision policy, when LossScaleOptimizer is used, loss is
      # scaled for numerical stability.
      if isinstance(
          optimizer, tf_keras.mixed_precision.LossScaleOptimizer):
        scaled_loss = optimizer.get_scaled_loss(scaled_loss)

    tvars = model.trainable_variables
    grads = tape.gradient(scaled_loss, tvars)
    # Scales back gradient before apply_gradients when LossScaleOptimizer is
    # used.
    if isinstance(optimizer, tf_keras.mixed_precision.LossScaleOptimizer):
      grads = optimizer.get_unscaled_gradients(grads)
    optimizer.apply_gradients(list(zip(grads, tvars)))

    logs = all_losses
    if metrics:
      self.process_metrics(metrics, labels, outputs)
      logs.update({m.name: m.result() for m in metrics})
    elif model.compiled_metrics:
      self.process_compiled_metrics(model.compiled_metrics, labels, outputs)
      logs.update({m.name: m.result() for m in model.metrics})
    return logs
//  [python]
def _eval_step(self, batch):
        features, input_length, phone_labels, phone_label_length, tar_label = batch
        max_length = tf.shape(tar_label)[1]
        enc_output = self.encoder(features, training=False)
        ctc_output = self.ctc_model(enc_output, training=False)

        ctc_output = tf.nn.softmax(ctc_output, -1)
        ctc_loss = tf.keras.backend.ctc_batch_cost(t.cast(phone_labels, tf.int32),
                                                   t.cast(ctc_output, tf.float32),
                                                   t.cast(input_length[:, tf.newaxis], tf.int32),
                                                   t.cast(phone_label_length[:, tf.newaxis], tf.int32),
                                                   )
        ctc_decode = tf.keras.backend.ctc_decode(ctc_output, input_length)[0][0]
        ctc_decode = t.cast(tf.clip_by_value(ctc_decode, 0, self.phone_featurizer.num_classes), tf.int32)

        ctc_out = self.translator([ctc_decode, enc_output], training=False)


        translate_loss = tf.keras.losses.sparse_categorical_crossentropy(tar_label, ctc_out[:, :max_length], True)
        ctc_acc = self.ctc_acc(phone_labels, ctc_decode)
        translate_acc = self.translate_acc(tar_label, ctc_out, max_length)
        self.eval_metrics["ctc_loss"].update_state(ctc_loss)
        self.eval_metrics["translate_loss"].update_state(translate_loss)
        self.eval_metrics['ctc_acc'].update_state(ctc_acc)
        self.eval_metrics['translate_acc'].update_state(translate_acc)