jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def _prune_weight(self, op, weight_tensor, scope):
        with tf.variable_scope(scope, reuse=tf.AUTO_REUSE):
            weight_variable = get_variable_by_tensor(weight_tensor)
            m = self._prune_var_mask_abs_grads[weight_variable]
            m_const = self._pname_sens_constants[weight_variable]
            m_sorted = tf.contrib.framework.sort(tf.reshape(m_const, [-1]))
            num = m_sorted.shape.num_elements() * self._config_dict[op.name] / self._total_pruning_iterations * tf.cast(self._step, tf.float32)
            num = tf.cast(num, tf.int64)
            threshold = m_sorted[num]
            mask = tf.cast(m_const >= threshold, tf.float32)
            self._mask_dict[op.name] = mask

            if self._debug_info:
                pruning_ratio = 1.0 - tf.reduce_mean(mask)
                print_op = tf.print("pruning iter:", self._step, "/", self._total_pruning_iterations, op.name, "pruning_ratio:", pruning_ratio)

                with tf.control_dependencies([print_op]):
                    pruned_weight_tensor = weight_variable * mask
            else:
                pruned_weight_tensor = weight_variable * mask

            new_weight_tensor_assign = tf.assign(weight_variable, pruned_weight_tensor)
            tf.add_to_collection("update_pruning_masks", new_weight_tensor_assign)

            return new_weight_tensor_assign
//  [python]
def _generate_random_positions(
      self, seed: Optional[int] = None) -> Tuple[tf.Tensor, tf.Tensor]:
    """Generates random instance positions in videos."""

    num_frames = self._num_frames * 2 if self._is_ssl else self._num_frames
    shape = [num_frames, self._num_instances, 4]
    xmin = tf.random.uniform(shape=shape[:-1],
                             minval=0.0,
                             maxval=1.0,
                             dtype=tf.float32,
                             seed=seed)
    ymin = tf.random.uniform(shape=shape[:-1],
                             minval=0.0,
                             maxval=1.0,
                             dtype=tf.float32,
                             seed=seed)
    xdelta = tf.random.uniform(shape=shape[:-1],
                               minval=0.1,
                               maxval=0.5,
                               dtype=tf.float32,
                               seed=seed)
    aspect_ratio = tf.random.uniform(shape=shape[:-1],
                                     minval=0.5,
                                     maxval=2.0,
                                     dtype=tf.float32,
                                     seed=seed)
    ydelta = xdelta * aspect_ratio
    xmax = tf.math.minimum(xmin + xdelta, 1.0 - 1e-3)
    ymax = tf.math.minimum(ymin + ydelta, 1.0 - 1e-3)
    random_positions = tf.stack([ymin, xmin, ymax, xmax], axis=-1)
    random_positions = tf.to_float(random_positions)
    instances_mask = tf.ones(shape[:-1], dtype=tf.bool)
    return random_positions, instances_mask
//  [python]
def _generate_random_positions(
      self, seed: Optional[int] = None) -> Tuple[tf.Tensor, tf.Tensor]:
    """Generates random instance positions in videos."""

    num_frames = self._num_frames * 2 if self._is_ssl else self._num_frames
    shape = [num_frames, self._num_instances, 4]
    xmin = tf.random.uniform(shape=shape[:-1],
                             minval=0.0,
                             maxval=1.0,
                             dtype=tf.float32,
                             seed=seed)
    ymin = tf.random.uniform(shape=shape[:-1],
                             minval=0.0,
                             maxval=1.0,
                             dtype=tf.float32,
                             seed=seed)
    xdelta = tf.random.uniform(shape=shape[:-1],
                               minval=0.1,
                               maxval=0.5,
                               dtype=tf.float32,
                               seed=seed)
    aspect_ratio = tf.random.uniform(shape=shape[:-1],
                                     minval=0.5,
                                     maxval=2.0,
                                     dtype=tf.float32,
                                     seed=seed)
    ydelta = xdelta * aspect_ratio
    xmax = tf.math.minimum(xmin + xdelta, 1.0 - 1e-3)
    ymax = tf.math.minimum(ymin + ydelta, 1.0 - 1e-3)
    random_positions = tf.stack([ymin, xmin, ymax, xmax], axis=-1)
    random_positions = tf.cast(random_positions, dtype=self._dtype)
    instances_mask = tf.ones(shape[:-1], dtype=tf.bool)
    return random_positions, instances_mask
//  [python]
def gradient_penalty(
        self, x: tf.Tensor, x_gen: tf.Tensor, labels: tf.Tensor
    ) -> tf.Tensor:
        """
        Compute gradient penalty: L2(grad - 1)^2.

        Args:
            x: input batch
            x_gen: generated images
            labels: labels associated with x (and thus with x_gen)

        Returns:
            penalty on discriminator gradient
        """

        epsilon = tf.random.uniform([tf.shape(x)[0], 1, 1, 1], 0.0, 1.0)
        x_hat = epsilon * x + (1 - epsilon) * x_gen
        labels = tf.cast(labels, tf.float32)
        with tf.GradientTape() as tape:
            tape.watch([x_hat, labels])
            d_hat = self.discriminator([x_hat, labels], training=True)
        gradients = tape.gradient(d_hat, x_hat)
        ddx = tf.sqrt(tf.reduce_sum(gradients ** 2, axis=[1, 2]))
        d_regularizer = tf.reduce_mean((ddx - 1.0) ** 2)
        return d_regularizer
//  [python]
def my_masked_softmax_cross_entropy(logits, labels, mask):
        """Softmax cross-entropy loss with masking."""
        logits = tf.clip_by_value(logits, clip_value_min=1e-5, clip_value_max=1.)
        loss = tf.nn.softmax_cross_entropy_with_logits(logits=logits, labels=labels)
        mask = tf.cast(mask, dtype=tf.float32)
        mask /= tf.reduce_mean(mask)
        loss *= mask
        return tf.reduce_mean(loss)
//  [python]
def padded_one_hot_encoding(indices, depth, left_pad):
  """Returns a zero padded one-hot tensor.

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

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

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

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

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

  rank = len(indices.get_shape().as_list())
  if rank != 1:
    raise ValueError('`indices` must have rank 1, but has rank=%s' % rank)

  def one_hot_and_pad():
    one_hot = tf.cast(tf.one_hot(tf.cast(indices, tf.int64), depth,
                                 on_value=1, off_value=0), tf.float32)
    return tf.pad(one_hot, [[0, 0], [left_pad, 0]], mode='CONSTANT')
  result = tf.cond(tf.greater(tf.size(indices), 0), one_hot_and_pad,
                   lambda: tf.zeros((depth + left_pad, 0)))
  return tf.reshape(result, [-1, depth + left_pad])
//  [python]
def masked_softmax_cross_entropy(logits, labels, mask):
        """Softmax cross-entropy loss with masking."""
        loss = tf.nn.softmax_cross_entropy_with_logits(logits=logits, labels=labels)
        mask = tf.to_float(mask)
        mask /= tf.reduce_mean(mask)
        loss *= mask
        return tf.reduce_mean(loss)
//  [python]
def masked_softmax_cross_entropy(logits, labels, mask):
        """Softmax cross-entropy loss with masking."""
        loss = tf.nn.softmax_cross_entropy_with_logits(logits=logits, labels=labels)
        mask = tf.cast(mask, dtype=tf.float32)
        mask /= tf.reduce_mean(mask)
        loss *= mask
        return tf.reduce_mean(loss)
//  [python]
def _attn(
        self,
        query: tf.Tensor,
        key: tf.Tensor,
        value: tf.Tensor,
        attention_mask: tf.Tensor | None = None,
        head_mask: tf.Tensor | None = None,
    ) -> Tuple[tf.Tensor, tf.Tensor]:
        # compute causal mask from causal mask buffer
        query_length, key_length = shape_list(query)[-2], shape_list(key)[-2]
        causal_mask = self.get_causal_mask(key_length, query_length)

        # Keep the attention weights computation in fp32 to avoid overflow issues
        query = tf.to_float(query)
        key = tf.to_float(key)

        attn_weights = tf.matmul(query, key, transpose_b=True)
        attn_weights = tf.where(causal_mask, attn_weights, self.get_masked_bias(attn_weights.dtype))

        attn_weights = attn_weights / self.scale_attn

        if attention_mask is not None:
            # Apply the attention mask
            attn_weights = attn_weights + attention_mask

        attn_weights = stable_softmax(attn_weights, axis=-1)
        attn_weights = tf.cast(attn_weights, value.dtype)
        attn_weights = self.attn_dropout(attn_weights)

        # Mask heads if we want to
        if head_mask is not None:
            attn_weights = attn_weights * head_mask

        attn_output = tf.matmul(attn_weights, value)

        return attn_output, attn_weights
//  [python]
def _attn(
        self,
        query: tf.Tensor,
        key: tf.Tensor,
        value: tf.Tensor,
        attention_mask: tf.Tensor | None = None,
        head_mask: tf.Tensor | None = None,
    ) -> Tuple[tf.Tensor, tf.Tensor]:
        # compute causal mask from causal mask buffer
        query_length, key_length = shape_list(query)[-2], shape_list(key)[-2]
        causal_mask = self.get_causal_mask(key_length, query_length)

        # Keep the attention weights computation in fp32 to avoid overflow issues
        query = tf.cast(query, tf.float32)
        key = tf.cast(key, tf.float32)

        attn_weights = tf.matmul(query, key, transpose_b=True)
        attn_weights = tf.where(causal_mask, attn_weights, self.get_masked_bias(attn_weights.dtype))

        attn_weights = attn_weights / self.scale_attn

        if attention_mask is not None:
            # Apply the attention mask
            attn_weights = attn_weights + attention_mask

        attn_weights = stable_softmax(attn_weights, axis=-1)
        attn_weights = tf.cast(attn_weights, value.dtype)
        attn_weights = self.attn_dropout(attn_weights)

        # Mask heads if we want to
        if head_mask is not None:
            attn_weights = attn_weights * head_mask

        attn_output = tf.matmul(attn_weights, value)

        return attn_output, attn_weights
//  [python]
def __init__(self,
               config: Optional[ModelConfig] = None,
               overrides: Optional[Dict[Text, Any]] = None):
    """Create a MobilenetEdgeTPU model.

    Args:
      config: (optional) the main model parameters to create the model
      overrides: (optional) a dict containing keys that can override config
    """
    overrides = overrides or {}
    config = config or ModelConfig()

    self.config = config.replace(**overrides)

    input_channels = self.config.input_channels
    model_name = self.config.model_name
    if isinstance(self.config.resolution, tuple):
      input_shape = (self.config.resolution[0], self.config.resolution[1],
                     input_channels)
    else:
      input_shape = (self.config.resolution, self.config.resolution,
                     input_channels)
    image_input = tf_keras.layers.Input(shape=input_shape)

    output = mobilenet_edgetpu_v1_model_blocks.mobilenet_edgetpu(
        image_input, self.config)

    if not isinstance(output, dict):
      # Cast to float32 in case we have a different model dtype
      output = tf.cast(output, tf.float32)
      self._output_specs = output.get_shape()
    else:
      self._output_specs = {
          feature: output[feature].get_shape() for feature in output
      }

    logging.info('Building model %s with params %s',
                 model_name,
                 self.config)

    super(MobilenetEdgeTPU, self).__init__(
        inputs=image_input, outputs=output, name=model_name)
//  [python]
def masked_sigmoid_cross_entropy(logits, labels, mask):
        """Softmax cross-entropy loss with masking."""
        labels = t.cast(labels, dtype=t.float32)
        loss = tf.nn.sigmoid_cross_entropy_with_logits(logits=logits, labels=labels)
        loss = tf.reduce_mean(loss, axis=1)
        mask = t.cast(mask, dtype=t.float32)
        mask /= tf.reduce_mean(mask)
        loss *= mask
        return tf.reduce_mean(loss)
//  [python]
def masked_sigmoid_cross_entropy(logits, labels, mask):
        """Softmax cross-entropy loss with masking."""
        labels = tf.cast(labels, dtype=tf.float32)
        loss = tf.nn.sigmoid_cross_entropy_with_logits(logits=logits, labels=labels)
        loss = tf.reduce_mean(loss, axis=1)
        mask = tf.cast(mask, dtype=tf.float32)
        mask /= tf.reduce_mean(mask)
        loss *= mask
        return tf.reduce_mean(loss)
//  [python]
def masked_accuracy(logits, labels, mask):
        """Accuracy with masking."""
        correct_prediction = tf.equal(tf.argmax(logits, 1), tf.argmax(labels, 1))
        accuracy_all = tf.to_float(correct_prediction)
        mask = tf.to_float(mask)
        mask /= tf.reduce_mean(mask)
        accuracy_all *= mask
        return tf.reduce_mean(accuracy_all)
//  [python]
def masked_accuracy(logits, labels, mask):
        """Accuracy with masking."""
        correct_prediction = tf.equal(tf.argmax(logits, 1), tf.argmax(labels, 1))
        accuracy_all = tf.cast(correct_prediction, tf.float32)
        mask = tf.cast(mask, dtype=tf.float32)
        mask /= tf.reduce_mean(mask)
        accuracy_all *= mask
        return tf.reduce_mean(accuracy_all)
//  [python]
def call(
        self,
        hidden_states: tf.Tensor,
        layer_past: Optional[Tuple[tf.Tensor, tf.Tensor]] = None,
        attention_mask: tf.Tensor | None = None,
        position_ids: tf.Tensor | None = None,
        head_mask: tf.Tensor | None = None,
        use_cache: bool = False,
        output_attentions: bool = False,
    ):
        query = self.q_proj(hidden_states)
        key = self.k_proj(hidden_states)
        value = self.v_proj(hidden_states)

        query = self._split_heads(query, True)
        key = self._split_heads(key, True)
        value = self._split_heads(value, False)

        sincos = tf.gather(self.embed_positions, position_ids, axis=0)
        sincos = tf.split(sincos, 2, axis=-1)
        if self.rotary_dim is not None:
            k_rot = key[:, :, :, : self.rotary_dim]
            k_pass = key[:, :, :, self.rotary_dim :]

            q_rot = query[:, :, :, : self.rotary_dim]
            q_pass = query[:, :, :, self.rotary_dim :]

            k_rot = apply_rotary_pos_emb(k_rot, sincos)
            q_rot = apply_rotary_pos_emb(q_rot, sincos)

            key = tf.concat((k_rot, k_pass), axis=-1)
            query = tf.concat((q_rot, q_pass), axis=-1)
        else:
            key = apply_rotary_pos_emb(key, sincos)
            query = apply_rotary_pos_emb(query, sincos)

        key = tf.transpose(key, (0, 2, 1, 3))
        query = tf.transpose(query, (0, 2, 1, 3))

        if layer_past is not None:
            past_key = layer_past[0]
            past_value = layer_past[1]
            key = tf.concat((past_key, key), axis=-2)
            value = tf.concat((past_value, value), axis=-2)

        if use_cache is True:
            present = (key, value)
        else:
            present = None

        # compute self-attention: V x Softmax(QK^T)
        attn_output, attn_weights = self._attn(query, key, value, attention_mask, head_mask)

        attn_output = self._merge_heads(attn_output)
        attn_output = self.out_proj(attn_output)
        attn_output = self.resid_dropout(attn_output)

        outputs = (attn_output, present)
        if output_attentions:
            outputs += (attn_weights,)

        return outputs
//  [python]
def call(
        self,
        hidden_states: tf.Tensor,
        layer_past: Optional[Tuple[tf.Tensor, tf.Tensor]] = None,
        attention_mask: tf.Tensor | None = None,
        position_ids: tf.Tensor | None = None,
        head_mask: tf.Tensor | None = None,
        use_cache: bool = False,
        output_attentions: bool = False,
    ):
        query = self.q_proj(hidden_states)
        key = self.k_proj(hidden_states)
        value = self.v_proj(hidden_states)

        query = self._split_heads(query, True)
        key = self._split_heads(key, True)
        value = self._split_heads(value, False)

        sincos = tf.cast(tf.gather(self.embed_positions, position_ids, axis=0), hidden_states.dtype)
        sincos = tf.split(sincos, 2, axis=-1)
        if self.rotary_dim is not None:
            k_rot = key[:, :, :, : self.rotary_dim]
            k_pass = key[:, :, :, self.rotary_dim :]

            q_rot = query[:, :, :, : self.rotary_dim]
            q_pass = query[:, :, :, self.rotary_dim :]

            k_rot = apply_rotary_pos_emb(k_rot, sincos)
            q_rot = apply_rotary_pos_emb(q_rot, sincos)

            key = tf.concat((k_rot, k_pass), axis=-1)
            query = tf.concat((q_rot, q_pass), axis=-1)
        else:
            key = apply_rotary_pos_emb(key, sincos)
            query = apply_rotary_pos_emb(query, sincos)

        key = tf.transpose(key, (0, 2, 1, 3))
        query = tf.transpose(query, (0, 2, 1, 3))

        if layer_past is not None:
            past_key = layer_past[0]
            past_value = layer_past[1]
            key = tf.concat((past_key, key), axis=-2)
            value = tf.concat((past_value, value), axis=-2)

        if use_cache is True:
            present = (key, value)
        else:
            present = None

        # compute self-attention: V x Softmax(QK^T)
        attn_output, attn_weights = self._attn(query, key, value, attention_mask, head_mask)

        attn_output = self._merge_heads(attn_output)
        attn_output = self.out_proj(attn_output)
        attn_output = self.resid_dropout(attn_output)

        outputs = (attn_output, present)
        if output_attentions:
            outputs += (attn_weights,)

        return outputs
//  [python]
def micro_f1(logits, labels, mask):
        """Accuracy with masking."""
        predicted = tf.round(tf.nn.sigmoid(logits))

        # Use integers to avoid any nasty FP behaviour
        predicted = tf.to_int32(predicted)
        labels = tf.to_int32(labels)
        mask = tf.to_int32(mask)

        # expand the mask so that broadcasting works ([nb_nodes, 1])
        mask = tf.expand_dims(mask, -1)

        # Count true positives, true negatives, false positives and false negatives.
        tp = tf.count_nonzero(predicted * labels * mask)
        tn = tf.count_nonzero((predicted - 1) * (labels - 1) * mask)
        fp = tf.count_nonzero(predicted * (labels - 1) * mask)
        fn = tf.count_nonzero((predicted - 1) * labels * mask)

        # Calculate accuracy, precision, recall and F1 score.
        precision = tp / (tp + fp)
        recall = tp / (tp + fn)
        fmeasure = (2 * precision * recall) / (precision + recall)
        fmeasure = tf.to_float(fmeasure)
        return fmeasure
//  [python]
def micro_f1(logits, labels, mask):
        """Accuracy with masking."""
        predicted = tf.round(tf.nn.sigmoid(logits))

        # Use integers to avoid any nasty FP behaviour
        predicted = tf.cast(predicted, dtype=tf.int32)
        labels = tf.cast(labels, dtype=tf.int32)
        mask = tf.cast(mask, dtype=tf.int32)

        # expand the mask so that broadcasting works ([nb_nodes, 1])
        mask = tf.expand_dims(mask, -1)

        # Count true positives, true negatives, false positives and false negatives.
        tp = tf.count_nonzero(predicted * labels * mask)
        tn = tf.count_nonzero((predicted - 1) * (labels - 1) * mask)
        fp = tf.count_nonzero(predicted * (labels - 1) * mask)
        fn = tf.count_nonzero((predicted - 1) * labels * mask)

        # Calculate accuracy, precision, recall and F1 score.
        precision = tp / (tp + fp)
        recall = tp / (tp + fn)
        fmeasure = (2 * precision * recall) / (precision + recall)
        fmeasure = tf.cast(fmeasure, tf.float32)
        return fmeasure
//  [python]
def maybe_apply_gradients(optimizer, gradients, step, update_ops=None):
    """If `gradients` is none, increase `step` by 1. Otherwise apply gradient as normal."""
    dependencies = []
    if update_ops is not None:
      dependencies = update_ops
    with tf.control_dependencies(dependencies):
      if gradients:
        return optimizer.apply_gradients(gradients, global_step=step)
      else:
        return tf.to_int32(step.assign(step + 1))
//  [python]
def maybe_apply_gradients(optimizer, gradients, step, update_ops=None):
    """If `gradients` is none, increase `step` by 1. Otherwise apply gradient as normal."""
    dependencies = []
    if update_ops is not None:
      dependencies = update_ops
    with tf.control_dependencies(dependencies):
      if gradients:
        return optimizer.apply_gradients(gradients, global_step=step)
      else:
        return tf.cast(step.assign(step + 1), tf.bool)
//  [python]
def _force_data_dependency(
    first_compute: Sequence[tf.Tensor],
    then_compute: Sequence[tf.Tensor]) -> List[tf.Tensor]:
  """Force all of `then_compute` to depend on all of `first_compute`.

  Uses a dummy data dependency, which is useful when running on TPUs because
  XLA ignores control dependencies. Only supports float arguments.

  Args:
    first_compute: Sequence of `Tensor`s to be executed before `then_compute`.
    then_compute: Sequence of `Tensor`s to executed after `first_compute`.

  Returns:
    Sequence of `Tensor`s with same length of `then_compute`.

  Raises:
    ValueError: if ranks are unknown or types are not floating.
  """

  def _first_element(x):
    if x.shape.ndims is None:
      raise ValueError('Rank of Tensor %s must be known' % x)
    ndims = x.shape.ndims
    begin = tf.zeros(ndims, dtype=tf.int32)
    size = tf.ones(ndims, dtype=tf.int32)
    return tf.reshape(tf.slice(x, begin, size), [])

  first_compute_sum = tf.add_n(
      [_first_element(x) for x in first_compute if x is not None])
  dtype = first_compute_sum.dtype
  if not dtype.is_floating:
    raise ValueError('_force_data_dependency only supports floating dtypes.')
  zero = np.finfo(dtype.as_numpy_dtype).tiny * first_compute_sum
  return [
      x + t.cast(zero, x.dtype) if x is not None else None
      for x in then_compute
  ]
//  [python]
def _force_data_dependency(
    first_compute: Sequence[tf.Tensor],
    then_compute: Sequence[tf.Tensor]) -> List[tf.Tensor]:
  """Force all of `then_compute` to depend on all of `first_compute`.

  Uses a dummy data dependency, which is useful when running on TPUs because
  XLA ignores control dependencies. Only supports float arguments.

  Args:
    first_compute: Sequence of `Tensor`s to be executed before `then_compute`.
    then_compute: Sequence of `Tensor`s to executed after `first_compute`.

  Returns:
    Sequence of `Tensor`s with same length of `then_compute`.

  Raises:
    ValueError: if ranks are unknown or types are not floating.
  """

  def _first_element(x):
    if x.shape.ndims is None:
      raise ValueError('Rank of Tensor %s must be known' % x)
    ndims = x.shape.ndims
    begin = tf.zeros(ndims, dtype=tf.int32)
    size = tf.ones(ndims, dtype=tf.int32)
    return tf.reshape(tf.slice(x, begin, size), [])

  first_compute_sum = tf.add_n(
      [_first_element(x) for x in first_compute if x is not None])
  dtype = first_compute_sum.dtype
  if not dtype.is_floating:
    raise ValueError('_force_data_dependency only supports floating dtypes.')
  zero = np.finfo(dtype.as_numpy_dtype).tiny * first_compute_sum
  return [
      x + tf.cast(zero, x.dtype) if x is not None else None
      for x in then_compute
  ]
//  [python]
def get_growing_source_and_target(data_batched, global_step):
    # TODO: should I grow the source as well?
    print('TODO: should I grow the source as well? Run an experiment on this!')
    sources = data_batched.get('source')
    targets = data_batched.get('target')
    if FLAGS.is_growing:
      with tf.variable_scope('alpha_grow'):
        alpha_grow = tf.to_float(global_step - FLAGS.grow_start_number_of_steps) / (
            FLAGS.max_number_of_steps - FLAGS.grow_start_number_of_steps)
        if sources is not None:
          sources = GanModel.get_growing_image(sources, alpha_grow, name_postfix='sources')
        if targets is not None:
          targets = GanModel.get_growing_image(targets, alpha_grow, name_postfix='targets')
    else:
      alpha_grow = 0.0
    return sources, targets, alpha_grow
//  [python]
def get_growing_source_and_target(data_batched, global_step):
    # TODO: should I grow the source as well?
    print('TODO: should I grow the source as well? Run an experiment on this!')
    sources = data_batched.get('source')
    targets = data_batched.get('target')
    if FLAGS.is_growing:
      with tf.variable_scope('alpha_grow'):
        alpha_grow = tf.cast(global_step - FLAGS.grow_start_number_of_steps, targets.dtype) / (
            FLAGS.max_number_of_steps - FLAGS.grow_start_number_of_steps)
        if sources is not None:
          sources = GanModel.get_growing_image(sources, alpha_grow, name_postfix='sources')
        if targets is not None:
          targets = GanModel.get_growing_image(targets, alpha_grow, name_postfix='targets')
    else:
      alpha_grow = 0.0
    return sources, targets, alpha_grow
//  [python]
def _add_wgan_gp_loss(real_image, generated_image, discriminator_network_fn, loss_scope='gradient_penalty',
                        discriminator_var_scope=DISCRIMINATOR_SCOPE, overall_weight=1.0):
    tf.logging.info('wgan_gp does not interact well when discriminator uses batch norm (according to wgan-gp paper).'
                    'Double check the discriminator architecture if necessary.')
    with tf.variable_scope('wgan_gp'):
      alpha_shape = tf.TensorShape(
        [real_image.shape[0]] + [1 for _ in range(1, len(real_image.shape))])
      alpha = tf.random_uniform(shape=alpha_shape, minval=0., maxval=1., name='alpha')
      real_image_casted = tf.cast(real_image, generated_image.dtype)
      interpolates = real_image_casted + alpha * (generated_image - real_image_casted)
    with tf.variable_scope(discriminator_var_scope, reuse=True):
      interpolate_prediction, interpolate_end_points = discriminator_network_fn(interpolates)
      interpolate_prediction = tf.cast(interpolate_prediction, tf.float32)

    interpolate_gradients = tf.gradients(interpolate_prediction, [interpolates])[0]
    interpolate_gradients = tf.cast(interpolate_gradients, tf.float32)
    interpolate_slopes = tf.sqrt(tf.reduce_sum(tf.square(interpolate_gradients), axis=[1, 2, 3]))

    gradient_penalty = tf.reduce_mean(tf.square((interpolate_slopes - 1.0)))
    # WGAN-GP loss
    gradient_penalty = tf.losses.compute_weighted_loss(gradient_penalty,
                                                       weights=FLAGS.gradient_penalty_lambda * overall_weight,
                                                       scope=loss_scope,
                                                       loss_collection=DISCRIMINATOR_LOSS_COLLECTION)
    return gradient_penalty
//  [python]
def _add_dragan_loss(real_image, discriminator_network_fn, loss_scope='gradient_penalty',
                       var_scope=DISCRIMINATOR_SCOPE,
                       overall_weight=1.0):
    tf.logging.info('Dragan does not interact well with batch norm in both generator and discriminator, according to '
                    '"How to Train Your DRAGAN". Please double check your network setup.')
    with tf.variable_scope('dragan'):
      alpha_shape = tf.TensorShape([real_image.shape[0]] + [1 for _ in range(1, len(real_image.shape))])
      alpha = tf.random_uniform(shape=alpha_shape, minval=0., maxval=1., name='alpha', dtype=real_image.dtype)
      difference = GanModel.get_perturbed_batch(real_image) - real_image
      interpolates = real_image + alpha * difference
    with tf.variable_scope(var_scope, reuse=True):
      interpolate_prediction, interpolate_end_points = discriminator_network_fn(interpolates)
      interpolate_prediction = tf.to_float(interpolate_prediction)

    interpolate_gradients = tf.gradients(interpolate_prediction, [interpolates])[0]
    interpolate_gradients = tf.to_float(interpolate_gradients)
    interpolate_slopes = tf.sqrt(tf.reduce_sum(tf.square(interpolate_gradients),
                                               reduction_indices=[i for i in
                                                                  range(1, len(interpolate_gradients.shape))]))
    gradient_penalty = tf.reduce_mean((interpolate_slopes - 1.0) ** 2)
    gradient_penalty = tf.losses.compute_weighted_loss(gradient_penalty,
                                                       weights=FLAGS.gradient_penalty_lambda * overall_weight,
                                                       scope=loss_scope,
                                                       loss_collection=DISCRIMINATOR_LOSS_COLLECTION)
    return gradient_penalty
//  [python]
def _add_dragan_loss(real_image, discriminator_network_fn, loss_scope='gradient_penalty',
                       var_scope=DISCRIMINATOR_SCOPE,
                       overall_weight=1.0):
    tf.logging.info('Dragan does not interact well with batch norm in both generator and discriminator, according to '
                    '"How to Train Your DRAGAN". Please double check your network setup.')
    with tf.variable_scope('dragan'):
      alpha_shape = tf.TensorShape([real_image.shape[0]] + [1 for _ in range(1, len(real_image.shape))])
      alpha = tf.random_uniform(shape=alpha_shape, minval=0., maxval=1., name='alpha', dtype=real_image.dtype)
      difference = GanModel.get_perturbed_batch(real_image) - real_image
      interpolates = real_image + alpha * difference
    with tf.variable_scope(var_scope, reuse=True):
      interpolate_prediction, interpolate_end_points = discriminator_network_fn(interpolates)
      interpolate_prediction = tf.cast(interpolate_prediction, tf.float32)

    interpolate_gradients = tf.gradients(interpolate_prediction, [interpolates])[0]
    interpolate_gradients = tf.cast(interpolate_gradients, tf.float32)
    interpolate_slopes = tf.sqrt(tf.reduce_sum(tf.square(interpolate_gradients),
                                               reduction_indices=[i for i in
                                                                  range(1, len(interpolate_gradients.shape))]))
    gradient_penalty = tf.reduce_mean((interpolate_slopes - 1.0) ** 2)
    gradient_penalty = tf.losses.compute_weighted_loss(gradient_penalty,
                                                       weights=FLAGS.gradient_penalty_lambda * overall_weight,
                                                       scope=loss_scope,
                                                       loss_collection=DISCRIMINATOR_LOSS_COLLECTION)
    return gradient_penalty
//  [python]
def call(self, alignments, mask=None):
    """Computes homology detection logits from SW scores and seq lengths.

    Args:
      alignments: a 2-tuple of scores and paths for the batch.
      mask: a single tf.Tensor<float>[batch, 2, len], corresponding to the
        paddings masks for the two sequences.

    Returns:
      A tf.Tensor<float>[batch, 1] with the logits for each example in the
      batch.
    """
    scores = alignments[0]
    length_fn = lambda x: tf.reduce_sum(x, axis=1)
    masks = t.cast(mask, tf.float32)
    mn = t.cast(length_fn(masks[:, 0]) * length_fn(masks[:, 1]), scores.dtype)
    logits = (self.b + tf.exp(self.log_l) * scores -
              tf.exp(self.log_k) * tf.math.log(mn))
    return logits[:, tf.newaxis]
//  [python]
def call(self, alignments, mask=None):
    """Computes homology detection logits from SW scores and seq lengths.

    Args:
      alignments: a 2-tuple of scores and paths for the batch.
      mask: a single tf.Tensor<float>[batch, 2, len], corresponding to the
        paddings masks for the two sequences.

    Returns:
      A tf.Tensor<float>[batch, 1] with the logits for each example in the
      batch.
    """
    scores = alignments[0]
    length_fn = lambda x: tf.reduce_sum(x, axis=1)
    masks = tf.cast(mask, tf.float32)
    mn = tf.cast(length_fn(masks[:, 0]) * length_fn(masks[:, 1]), scores.dtype)
    logits = (self.b + tf.exp(self.log_l) * scores -
              tf.exp(self.log_k) * tf.math.log(mn))
    return logits[:, tf.newaxis]
//  [python]
def call(self, alignments, mask=None):
    """Computes homology detection logits from SW scores and seq lengths.

    Args:
      alignments: a 2-tuple of scores and paths for the batch.
      mask: a single tf.Tensor<float>[batch, 2, len], corresponding to the
        paddings masks for the two sequences.

    Returns:
      A tf.Tensor<float>[batch, 1] with the logits for each example in the
      batch.
    """
    scores = alignments[0]
    masks = t.cast(mask, t.float32)
    length_fn = lambda x: t.math.reduce_sum(x, axis=1)
    bias = self.b + self.log_k2 - self.log_k1 + self.log_l2 - self.log_l1
    lin = (t.exp(self.log_l1) - t.exp(self.log_l2)) * scores
    mn = t.cast(length_fn(masks[:, 0]) * length_fn(masks[:, 1]), scores.dtype)
    exp = mn * (t.exp(self.log_k2 - t.exp(self.log_l2) * scores) -
                t.exp(self.log_k1 - t.exp(self.log_l1) * scores))
    logits = bias + lin - exp
    return logits[:, t.newaxis]
//  [python]
def call(self, alignments, mask=None):
    """Computes homology detection logits from SW scores and seq lengths.

    Args:
      alignments: a 2-tuple of scores and paths for the batch.
      mask: a single tf.Tensor<float>[batch, 2, len], corresponding to the
        paddings masks for the two sequences.

    Returns:
      A tf.Tensor<float>[batch, 1] with the logits for each example in the
      batch.
    """
    scores = alignments[0]
    masks = tf.cast(mask, tf.float32)
    length_fn = lambda x: tf.math.reduce_sum(x, axis=1)
    bias = self.b + self.log_k2 - self.log_k1 + self.log_l2 - self.log_l1
    lin = (tf.exp(self.log_l1) - tf.exp(self.log_l2)) * scores
    mn = tf.cast(length_fn(masks[:, 0]) * length_fn(masks[:, 1]), scores.dtype)
    exp = mn * (tf.exp(self.log_k2 - tf.exp(self.log_l2) * scores) -
                tf.exp(self.log_k1 - tf.exp(self.log_l1) * scores))
    logits = bias + lin - exp
    return logits[:, tf.newaxis]
//  [python]
def resize_axis(tensor, axis, new_size, fill_value=0):
  """Truncates or pads a tensor to new_size on a given axis.

  Truncate or extend tensor such that tensor.shape[axis] == new_size. If the
  size increases, the padding will be performed at the end, using fill_value.

  Args:
    tensor: The tensor to be resized.
    axis: An integer representing the dimension to be sliced.
    new_size: An integer or 0d tensor representing the new value for
      tensor.shape[axis].
    fill_value: Value to use to fill any new entries in the tensor. Will be cast
      to the type of tensor.

  Returns:
    The resized tensor.
  """
  tensor = tf.convert_to_tensor(tensor)
  shape = tf.unstack(tf.shape(tensor))

  pad_shape = shape[:]
  pad_shape[axis] = tf.maximum(0, new_size - shape[axis])

  shape[axis] = tf.minimum(shape[axis], new_size)
  shape = tf.stack(shape)

  resized = tf.concat([
      tf.slice(tensor, tf.zeros_like(shape), shape),
      tf.fill(tf.stack(pad_shape), tf.cast(fill_value, tensor.dtype))
  ], axis)

  # Update shape.
  new_shape = tensor.shape.as_list()  # A copy is being made.
  new_shape[axis] = new_size
  resized = tf.ensure_shape(resized, new_shape)
  return resized
//  [python]
def decode(self,
             serialized_example: tf.train.SequenceExample) -> Dict[str, Any]:
    """Parses a single tf.train.SequenceExample into video and label tensors."""
    contexts, features = tf.io.parse_single_sequence_example(
        serialized_example,
        context_features=self._context_features,
        sequence_features=self._sequence_features)
    decoded_tensor = {**contexts, **features}
    for i, name in enumerate(self._feature_names):
      # Convert the VarLen feature to dense tensor.
      if self._feature_from_bytes[i]:
        dtype = tf.dtypes.as_dtype(self._feature_dtypes[i])
        decoded_tensor[name] = tf.to_float(
            tf.io.decode_raw(decoded_tensor[name], out_type=dtype)
        )
      else:
        if isinstance(decoded_tensor[name], tf.SparseTensor):
          decoded_tensor[name] = tf.sparse.to_dense(decoded_tensor[name])
    return decoded_tensor
//  [python]
def decode(self,
             serialized_example: tf.train.SequenceExample) -> Dict[str, Any]:
    """Parses a single tf.train.SequenceExample into video and label tensors."""
    contexts, features = tf.io.parse_single_sequence_example(
        serialized_example,
        context_features=self._context_features,
        sequence_features=self._sequence_features)
    decoded_tensor = {**contexts, **features}
    for i, name in enumerate(self._feature_names):
      # Convert the VarLen feature to dense tensor.
      if self._feature_from_bytes[i]:
        dtype = tf.dtypes.as_dtype(self._feature_dtypes[i])
        decoded_tensor[name] = tf.cast(
            tf.io.decode_raw(decoded_tensor[name], out_type=dtype), tf.float32
        )
      else:
        if isinstance(decoded_tensor[name], tf.SparseTensor):
          decoded_tensor[name] = tf.sparse.to_dense(decoded_tensor[name])
    return decoded_tensor
//  [python]
def read_tensor_from_byte_array(byteArray,
                                input_height=224,
                                input_width=224,
                                input_mean=127.5,
                                input_std=127.5):
  input_name = "file_reader"
  output_name = "normalized"
  image_reader = tf.image.decode_image(byteArray,channels=3)
  float_caster = tf.to_float(image_reader)
  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 read_tensor_from_byte_array(byteArray,
                                input_height=224,
                                input_width=224,
                                input_mean=127.5,
                                input_std=127.5):
  input_name = "file_reader"
  output_name = "normalized"
  image_reader = tf.image.decode_image(byteArray,channels=3)
  float_caster = tf.cast(image_reader, 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 _reshape_and_cast(self, t):
    return tf.reshape(t, [-1]).astype(tf.int32)
//  [python]
def _reshape_and_cast(self, t):
    return tf.cast(tf.reshape(t, [-1]), tf.int32)
//  [python]
def pad_to_size(input_tensor, size):
  """Pads data with zeros to a given length at the first dimension if needed.

  Args:
    input_tensor: `Tensor` with any dimension.
    size: `int` number for the first dimension of output Tensor.

  Returns:
    `Tensor` with the first dimension padded to `size` if the first diemsion
    is less than `size`, otherwise no padding.
  """
  input_shape = tf.shape(input_tensor)
  padding_shape = []

  # Computes the padding length on the first dimension.
  padding_length = tf.maximum(0, size - tf.shape(input_tensor)[0])
  assert_length = tf.Assert(
      tf.greater_equal(padding_length, 0), [padding_length])
  with tf.control_dependencies([assert_length]):
    padding_shape.append(padding_length)

  # Copies shapes of the rest of input shape dimensions.
  for i in range(1, len(input_shape)):
    padding_shape.append(tf.shape(input=input_tensor)[i])

  # Pads input tensor to the fixed first dimension.
  paddings = tf.zeros(padding_shape, dtype=input_tensor.dtype)
  padded_tensor = tf.concat([input_tensor, paddings], axis=0)
  return padded_tensor
//  [python]
def pad_to_size(input_tensor, size):
  """Pads data with zeros to a given length at the first dimension if needed.

  Args:
    input_tensor: `Tensor` with any dimension.
    size: `int` number for the first dimension of output Tensor.

  Returns:
    `Tensor` with the first dimension padded to `size` if the first diemsion
    is less than `size`, otherwise no padding.
  """
  input_shape = tf.shape(input_tensor)
  padding_shape = []

  # Computes the padding length on the first dimension.
  padding_length = tf.maximum(0, size - tf.shape(input_tensor)[0])
  assert_length = tf.Assert(
      tf.greater_equal(padding_length, 0), [padding_length])
  with tf.control_dependencies([assert_length]):
    padding_shape.append(padding_length)

  # Copies shapes of the rest of input shape dimensions.
  for i in range(1, len(input_shape)):
    padding_shape.append(tf.shape(input=input_tensor)[i])

  # Pads input tensor to the fixed first dimension.
  paddings = tf.cast(tf.zeros(padding_shape), input_tensor.dtype)
  padded_tensor = tf.concat([input_tensor, paddings], axis=0)
  return padded_tensor
//  [python]
def call(self, inputs):
    if not isinstance(self.kernel, random_variable.RandomVariable):
      return super().call(inputs)
    self.call_weights()
    input_shape = tf.shape(inputs)
    output_shape = tf.concat([input_shape[:-1], [self.units]], 0)
    sign_input = tf.to_int32(2 * tf.random.uniform(input_shape,
                                               minval=0,
                                               maxval=2,
                                               dtype=tf.int32) - 1)
                         
    sign_output = tf.to_float(2 * tf.random.uniform(output_shape,
                                                minval=0,
                                                maxval=2,
                                                dtype=inputs.dtype) - 1)
                          
    kernel_mean = self.kernel.distribution.mean()
    perturbation = self.kernel - kernel_mean
    if inputs.shape.ndims <= 2:
      outputs = tf.matmul(inputs, kernel_mean)
      outputs += tf.matmul(inputs * sign_input, perturbation) * sign_output
    else:
      outputs = tf.tensordot(inputs, kernel_mean, [[-1], [0]])
      outputs += tf.tensordot(inputs * sign_input,
                              perturbation,
                              [[-1], [0]]) * sign_output
    if self.use_bias:
      outputs = tf.nn.bias_add(outputs, self.bias)
    if self.activation is not None:
      outputs = self.activation(outputs)
    return outputs
//  [python]
def call(self, inputs):
    if not isinstance(self.kernel, random_variable.RandomVariable):
      return super().call(inputs)
    self.call_weights()
    input_shape = tf.shape(inputs)
    output_shape = tf.concat([input_shape[:-1], [self.units]], 0)
    sign_input = tf.cast(2 * tf.random.uniform(input_shape,
                                               minval=0,
                                               maxval=2,
                                               dtype=tf.int32) - 1,
                         inputs.dtype)
    sign_output = tf.cast(2 * tf.random.uniform(output_shape,
                                                minval=0,
                                                maxval=2,
                                                dtype=inputs.dtype) - 1,
                          inputs.dtype)
    kernel_mean = self.kernel.distribution.mean()
    perturbation = self.kernel - kernel_mean
    if inputs.shape.ndims <= 2:
      outputs = tf.matmul(inputs, kernel_mean)
      outputs += tf.matmul(inputs * sign_input, perturbation) * sign_output
    else:
      outputs = tf.tensordot(inputs, kernel_mean, [[-1], [0]])
      outputs += tf.tensordot(inputs * sign_input,
                              perturbation,
                              [[-1], [0]]) * sign_output
    if self.use_bias:
      outputs = tf.nn.bias_add(outputs, self.bias)
    if self.activation is not None:
      outputs = self.activation(outputs)
    return outputs
//  [python]
def call(self, inputs):
    if inputs.dtype != self.compute_dtype:
      inputs = t.cast(inputs, dtype=self.compute_dtype)

    batch_size = tf.shape(inputs)[0]
    input_dim = self.alpha.shape[-1]
    examples_per_model = batch_size // self.ensemble_size

    # TODO(ywenxu): Merge the following two cases.
    if self.rank > 1:
      input_dim = self.alpha.shape[-1]
      alpha = tf.reshape(tf.tile(self.alpha, [1, 1, examples_per_model]),
                         [self.rank, batch_size, input_dim])
      gamma = tf.reshape(tf.tile(self.gamma, [1, 1, examples_per_model]),
                         [self.rank, batch_size, self.units])
      perturb_inputs = tf.expand_dims(inputs, 0) * alpha
      outputs = super().call(perturb_inputs)
      outputs = tf.reduce_sum(outputs * gamma, axis=0)
      outputs = tf.reshape(
          outputs, [self.ensemble_size, examples_per_model, -1])
    else:
      inputs = tf.reshape(
          inputs, [self.ensemble_size, examples_per_model, input_dim])
      alpha = tf.expand_dims(self.alpha, 1)
      gamma = tf.expand_dims(self.gamma, 1)
      perturb_inputs = inputs * alpha
      outputs = super().call(perturb_inputs) * gamma
    if self.use_ensemble_bias:
      bias = tf.expand_dims(self.ensemble_bias, 1)
      outputs += bias
    if self.ensemble_activation is not None:
      outputs = self.ensemble_activation(outputs)
    outputs = tf.reshape(outputs, [batch_size, self.units])
    return outputs
//  [python]
def call(self, inputs):
    if inputs.dtype != self.compute_dtype:
      inputs = tf.cast(inputs, dtype=self.compute_dtype)

    batch_size = tf.shape(inputs)[0]
    input_dim = self.alpha.shape[-1]
    examples_per_model = batch_size // self.ensemble_size

    # TODO(ywenxu): Merge the following two cases.
    if self.rank > 1:
      input_dim = self.alpha.shape[-1]
      alpha = tf.reshape(tf.tile(self.alpha, [1, 1, examples_per_model]),
                         [self.rank, batch_size, input_dim])
      gamma = tf.reshape(tf.tile(self.gamma, [1, 1, examples_per_model]),
                         [self.rank, batch_size, self.units])
      perturb_inputs = tf.expand_dims(inputs, 0) * alpha
      outputs = super().call(perturb_inputs)
      outputs = tf.reduce_sum(outputs * gamma, axis=0)
      outputs = tf.reshape(
          outputs, [self.ensemble_size, examples_per_model, -1])
    else:
      inputs = tf.reshape(
          inputs, [self.ensemble_size, examples_per_model, input_dim])
      alpha = tf.expand_dims(self.alpha, 1)
      gamma = tf.expand_dims(self.gamma, 1)
      perturb_inputs = inputs * alpha
      outputs = super().call(perturb_inputs) * gamma
    if self.use_ensemble_bias:
      bias = tf.expand_dims(self.ensemble_bias, 1)
      outputs += bias
    if self.ensemble_activation is not None:
      outputs = self.ensemble_activation(outputs)
    outputs = tf.reshape(outputs, [batch_size, self.units])
    return outputs
//  [python]
def retrieve_seq_length_op2(data):
    """An op to compute the length of a sequence, from input shape of [batch_size, n_step(max)],
    it can be used when the features of padding (on right hand side) are all zeros.

    Parameters
    -----------
    data : tensor
        [batch_size, n_step(max)] with zero padding on right hand side.

    Examples
    --------
    >>> data = [[1,2,0,0,0],
    ...         [1,2,3,0,0],
    ...         [1,2,6,1,0]]
    >>> o = retrieve_seq_length_op2(data)
    >>> sess = tf.InteractiveSession()
    >>> tl.layers.initialize_global_variables(sess)
    >>> print(o.eval())
    ... [2 3 4]
    """
    return tf.reduce_sum(tf.cast(tf.greater(data, tf.zeros_like(data)), tf.int32), 1)