<commit_before># Copyright 2020 The TensorFlow Authors
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Module with loss functions."""

import collections
from typing import Sequence, Union

import tensorflow as tf


def gradient_penalty_loss(real_data: Union[tf.Tensor, Sequence[tf.Tensor]],
                          generated_data: Union[tf.Tensor, Sequence[tf.Tensor]],
                          discriminator: tf.keras.Model,
                          weight: float = 10.0,
                          eps: float = 1e-8,
                          name_scope: str = 'gradient_penalty') -> tf.Tensor:
  """Gradient penalty loss.

  This function implements the gradient penalty loss from "Improved Training of
  Wasserstein GANs"
  https://arxiv.org/abs/1704.00028
  This version also implements the multi-scale extension proposed in "MSG-GAN:
  Multi-Scale Gradient GAN for Stable Image Synthesis"
  https://arxiv.org/abs/1903.06048

  Args:
    real_data: Samples from the real data.
    generated_data: Samples from the generated data.
    discriminator: The Keras model of the discriminator. This model is expected
      to take as input a tf.Tensor or a sequence of tf.Tensor depending on what
      is provided as generated_data and real_data.
    weight: The weight of the loss.
    eps: A small positive value that is added to the argument of the square root
      to avoid undefined gradients.
    name_scope: The name scope of the loss.

  Returns:
    The gradient penalty loss.

  Raises:
    TypeError if real_data and generated_data are not both either tf.Tensor or
    both a sequence of tf.Tensor.
    ValueError if the numnber of elements in real_data and generated_data are
    not equal.
  """
  with tf.name_scope(name=name_scope):
    with tf.GradientTape() as tape:
      if (isinstance(real_data, tf.Tensor) and
          isinstance(generated_data, tf.Tensor)):
        epsilon = tf.random.uniform(
            [tf.shape(real_data)[0]] + [1] * (real_data.shape.ndims - 1),
            minval=0.0,
            maxval=1.0,
            dtype=real_data.dtype)
        interpolated_data = epsilon * real_data + (1.0 -
                                                   epsilon) * generated_data
      elif (isinstance(real_data, collections.Sequence) and
            isinstance(generated_data, collections.Sequence)):
        if len(real_data) != len(generated_data):
          raise ValueError(
              'The number of elements in real_data and generated_data are '
              'expected to be equal but got: %d and %d' %
              (len(real_data), len(generated_data)))
        epsilon = tf.random.uniform(
            [tf.shape(real_data[0])[0]] + [1] * (real_data[0].shape.ndims - 1),
            minval=0.0,
            maxval=1.0,
            dtype=real_data[0].dtype)
        interpolated_data = [
            epsilon * real_level + (1.0 - epsilon) * generated_level
            for real_level, generated_level in zip(real_data, generated_data)
        ]
      else:
        raise TypeError(
            'real_data and generated data should either both be a tf.Tensor '
            'or both a sequence of tf.Tensor but got: %s and %s' %
            (type(real_data), type(generated_data)))
      # By default the gradient tape only watches trainable variables.
      tape.watch(interpolated_data)<commit_msg>Make this code compatible with Python 3.10.<commit_after># Copyright 2020 The TensorFlow Authors
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Module with loss functions."""

from collections import abc
from typing import Sequence, Union

import tensorflow as tf


def gradient_penalty_loss(real_data: Union[tf.Tensor, Sequence[tf.Tensor]],
                          generated_data: Union[tf.Tensor, Sequence[tf.Tensor]],
                          discriminator: tf.keras.Model,
                          weight: float = 10.0,
                          eps: float = 1e-8,
                          name_scope: str = 'gradient_penalty') -> tf.Tensor:
  """Gradient penalty loss.

  This function implements the gradient penalty loss from "Improved Training of
  Wasserstein GANs"
  https://arxiv.org/abs/1704.00028
  This version also implements the multi-scale extension proposed in "MSG-GAN:
  Multi-Scale Gradient GAN for Stable Image Synthesis"
  https://arxiv.org/abs/1903.06048

  Args:
    real_data: Samples from the real data.
    generated_data: Samples from the generated data.
    discriminator: The Keras model of the discriminator. This model is expected
      to take as input a tf.Tensor or a sequence of tf.Tensor depending on what
      is provided as generated_data and real_data.
    weight: The weight of the loss.
    eps: A small positive value that is added to the argument of the square root
      to avoid undefined gradients.
    name_scope: The name scope of the loss.

  Returns:
    The gradient penalty loss.

  Raises:
    TypeError if real_data and generated_data are not both either tf.Tensor or
    both a sequence of tf.Tensor.
    ValueError if the numnber of elements in real_data and generated_data are
    not equal.
  """
  with tf.name_scope(name=name_scope):
    with tf.GradientTape() as tape:
      if (isinstance(real_data, tf.Tensor) and
          isinstance(generated_data, tf.Tensor)):
        epsilon = tf.random.uniform(
            [tf.shape(real_data)[0]] + [1] * (real_data.shape.ndims - 1),
            minval=0.0,
            maxval=1.0,
            dtype=real_data.dtype)
        interpolated_data = epsilon * real_data + (1.0 -
                                                   epsilon) * generated_data
      elif (isinstance(real_data, abc.Sequence) and
            isinstance(generated_data, abc.Sequence)):
        if len(real_data) != len(generated_data):
          raise ValueError(
              'The number of elements in real_data and generated_data are '
              'expected to be equal but got: %d and %d' %
              (len(real_data), len(generated_data)))
        epsilon = tf.random.uniform(
            [tf.shape(real_data[0])[0]] + [1] * (real_data[0].shape.ndims - 1),
            minval=0.0,
            maxval=1.0,
            dtype=real_data[0].dtype)
        interpolated_data = [
            epsilon * real_level + (1.0 - epsilon) * generated_level
            for real_level, generated_level in zip(real_data, generated_data)
        ]
      else:
        raise TypeError(
            'real_data and generated data should either both be a tf.Tensor '
            'or both a sequence of tf.Tensor but got: %s and %s' %
            (type(real_data), type(generated_data)))
      # By default the gradient tape only watches trainable variables.
      tape.watch(interpolated_data)
{'fraction_non_alphanumeric': 0.08012446518864255, 'fraction_numerical': 0.013872682484117724, 'mean_word_length': 2.4763406940063093, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 5, 'https://': 6, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2482739', 'n_tokens_mistral': 2175, 'n_tokens_neox': 2100, 'n_words': 765}