text: # coding=utf-8
  # Copyright 2022 The TensorFlow GAN 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
  #
  #     http://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.
  
  # Lint as: python2, python3
  """Keras-like layers and utilities that implement Spectral Normalization.
  
  Based on "Spectral Normalization for Generative Adversarial Networks" by Miyato,
  et al in ICLR 2018. https://openreview.net/pdf?id=B1QRgziT-
  """
  
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  import numbers
  import re
  
  from six.moves import range
  import tensorflow.compat.v1 as tf
  
  __all__ = [
      'compute_spectral_norm',
      'spectral_normalize',
      'spectral_norm_regularizer',
      'spectral_normalization_custom_getter',
  ]
  
  # tf.bfloat16 should work, but tf.matmul converts those to tf.float32 which then
  # can't directly be assigned back to the tf.bfloat16 variable.
  _OK_DTYPES_FOR_SPECTRAL_NORM = (tf.float16, tf.float32, tf.float64)
  _PERSISTED_U_VARIABLE_SUFFIX = 'spectral_norm_u'
  
  
  def compute_spectral_norm(w_tensor, power_iteration_rounds=1,
                            training=True, name=None):
    """Estimates the largest singular value in the weight tensor.
  
    **NOTE**: When `training=True`, repeatedly running inference actually changes
    the variables, since the spectral norm is repeatedly approximated by a power
    iteration method.
  
    Args:
      w_tensor: The weight matrix whose spectral norm should be computed.
      power_iteration_rounds: The number of iterations of the power method to
        perform. A higher number yields a better approximation.
      training: Whether to update the spectral normalization on variable
        access. This is useful to turn off during eval, for example, to not affect
        the graph during evaluation.
      name: An optional scope name.
  
    Returns:
      The largest singular value (the spectral norm) of w.
  
    Raises:
      ValueError: If TF is executing eagerly.
      ValueError: If called within a distribution strategy that is not supported.
    """
    if tf.executing_eagerly():
      # Under eager mode, get_variable() creates a new variable on every call.
      raise ValueError(
          '`compute_spectral_norm` doesn\'t work when executing eagerly.')
    with tf.variable_scope(name, 'spectral_norm'):
      # The paper says to flatten convnet kernel weights from
      # (C_out, C_in, KH, KW) to (C_out, C_in * KH * KW). But TensorFlow's Conv2D
      # kernel weight shape is (KH, KW, C_in, C_out), so it should be reshaped to
      # (KH * KW * C_in, C_out), and similarly for other layers that put output
      # channels as last dimension.
      # n.b. this means that w here is equivalent to w.T in the paper.
      w = tf.reshape(w_tensor, (-1, w_tensor.get_shape()[-1]))
  
      # Persisted approximation of first left singular vector of matrix `w`.
      # Requires an appropriate aggregation method since we explicitly control
      # updates.
      replica_context = tf.distribute.get_replica_context()
      if replica_context is None:  # cross repica strategy.
        # TODO(joelshor): Determine appropriate aggregation method.
        raise ValueError("spectral norm isn't supported in cross-replica "
                         "distribution strategy.")
      elif not tf.distribute.has_strategy():  # default strategy.
        aggregation = None
      else:
        aggregation = tf.VariableAggregation.ONLY_FIRST_REPLICA
  
      u_var = tf.get_variable(
          _PERSISTED_U_VARIABLE_SUFFIX,
          shape=(w.shape[0], 1),
          dtype=w.dtype,
          initializer=tf.initializers.random_normal(),
          trainable=False,
          aggregation=aggregation)
      u = u_var
  
      # Use power iteration method to approximate spectral norm.
      for _ in range(power_iteration_rounds):
        # `v` approximates the first right singular vector of matrix `w`.
        v = tf.nn.l2_normalize(tf.matmul(a=w, b=u, transpose_a=True))
        u = tf.nn.l2_normalize(tf.matmul(w, v))
  
      # Update persisted approximation.
      if training:
        with tf.control_dependencies([u_var.assign(u, name='update_u')]):
          u = tf.identity(u)
  
      u = tf.stop_gradient(u)
      v = tf.stop_gradient(v)
  
      # Largest singular value of `w`.
      spectral_norm = tf.matmul(tf.matmul(a=u, b=w, transpose_a=True), v)
      spectral_norm.shape.assert_is_fully_defined()
      spectral_norm.shape.assert_is_compatible_with([1, 1])
  
      return spectral_norm[0][0]
  
  
  def spectral_normalize(w,
                         power_iteration_rounds=1,
                         equality_constrained=True,
                         training=True,
                         name=None):
    """Normalizes a weight matrix by its spectral norm.
  
    **NOTE**: When `training=True`, repeatedly running inference actually changes
    the variables, since the spectral norm is repeatedly approximated by a power
    iteration method.
  
    Args:
      w: The weight matrix to be normalized.
      power_iteration_rounds: The number of iterations of the power method to
        perform. A higher number yields a better approximation.
      equality_constrained: If set to `True` will normalize the matrix such that
        its spectral norm is equal to 1, otherwise, will normalize the matrix such
        that its norm is at most 1.
      training: Whether to update the spectral normalization on variable
        access. This is useful to turn off during eval, for example, to not affect
        the graph during evaluation.
      name: An optional scope name.
  
    Returns:
      The input weight matrix, normalized so that its spectral norm is at most
      one.
    """
    with tf.variable_scope(name, 'spectral_normalize'):
      normalization_factor = compute_spectral_norm(
          w, power_iteration_rounds=power_iteration_rounds, training=training)
      if not equality_constrained:
        normalization_factor = tf.maximum(1., normalization_factor)
      w_normalized = w / normalization_factor
      return tf.reshape(w_normalized, w.get_shape())
  
  
  def spectral_norm_regularizer(scale, power_iteration_rounds=1,
                                training=True, scope=None):
    """Returns a function that can be used to apply spectral norm regularization.
  
    Small spectral norms enforce a small Lipschitz constant, which is necessary
    for Wasserstein GANs.
  
    **NOTE**: Repeatedly running inference actually changes the variables, since
    the spectral norm is repeatedly approximated by a power iteration method.
  
    Args:
      scale: A scalar multiplier. 0.0 disables the regularizer.
      power_iteration_rounds: The number of iterations of the power method to
        perform. A higher number yields a better approximation.
      training: Whether to update the spectral normalization on variable
        access. This is useful to turn off during eval, for example, to not affect
        the graph during evaluation.
      scope: An optional scope name.
  
    Returns:
      A function with the signature `sn(weights)` that applies spectral norm
      regularization.
  
    Raises:
      ValueError: If scale is negative or if scale is not a float.
    """
    if isinstance(scale, numbers.Integral):
      raise ValueError('scale cannot be an integer: %s' % scale)
    if isinstance(scale, numbers.Real):
      if scale < 0.0:
        raise ValueError(
            'Setting a scale less than 0 on a regularizer: %g' % scale)
      if scale == 0.0:
        tf.logging.info('Scale of 0 disables regularizer.')
        return lambda _: None
  
    def sn(weights, name=None):
      """Applies spectral norm regularization to weights."""
      with tf.name_scope(scope, 'SpectralNormRegularizer', [weights]) as name:
        scale_t = tf.convert_to_tensor(
            value=scale, dtype=weights.dtype.base_dtype, name='scale')
        return tf.multiply(
            scale_t,
            compute_spectral_norm(
                weights, power_iteration_rounds=power_iteration_rounds,
                training=training),
            name=name)
  
    return sn
  
  
  def _default_name_filter(name):
    """A filter function to identify common names of weight variables.
  
    Args:
      name: The variable name.
  
    Returns:
      Whether `name` is a standard name for a weight/kernel variables used in the
      Keras, tf.layers, tf.contrib.layers or tf.contrib.slim libraries.
    """
    match = re.match(r'(.*\/)?(depthwise_|pointwise_)?(weights|kernel)$', name)
    return match is not None
  
  
  def spectral_normalization_custom_getter(name_filter=_default_name_filter,
                                           power_iteration_rounds=1,
                                           equality_constrained=True,
                                           training=True):
    """Custom getter that performs Spectral Normalization on a weight tensor.
  
    Specifically it divides the weight tensor by its largest singular value. This
    is intended to stabilize GAN training, by making the discriminator satisfy a
    local 1-Lipschitz constraint.
  
    Based on [Spectral Normalization for Generative Adversarial Networks][sn-gan].
  
    [sn-gan]: https://openreview.net/forum?id=B1QRgziT-
  
    To reproduce an SN-GAN, apply this custom_getter to every weight tensor of
    your discriminator. The last dimension of the weight tensor must be the number
    of output channels.
  
    Apply this to layers by supplying this as the `custom_getter` of a
    `tf.variable_scope`. For example:
  
      with tf.variable_scope(
          'discriminator', custom_getter=spectral_normalization_custom_getter()):
        net = discriminator_fn(net)
  
    It is important to carefully select to which weights you want to apply
    Spectral Normalization. In general you want to normalize the kernels of
    convolution and dense layers, but you do not want to normalize biases. You
    also want to avoid normalizing batch normalization (and similar) variables,
    but in general such layers play poorly with Spectral Normalization, since the
    gamma can cancel out the normalization in other layers. By default we supply a
    filter that matches the kernel variable names of the dense and convolution
    layers of the tf.layers, tf.contrib.layers, tf.keras and tf.contrib.slim
    libraries. If you are using anything else you'll need a custom `name_filter`.
  
    This custom getter internally creates a variable used to compute the spectral
    norm by power iteration. It will update every time the variable is accessed,
    which means the normalized discriminator weights may change slightly whilst
    training the generator. Whilst unusual, this matches how the paper's authors
    implement it, and in general additional rounds of power iteration can't hurt.
  
    IMPORTANT: Keras does not respect the custom_getter supplied by the
    VariableScope, so this approach won't work for Keras layers. For Keras layers
    each layer needs to have Spectral Normalization explicitly applied. This
    should be accomplished using code like:
  
      my_layer = tf.keras.layers.SomeLayer()
      layer.build(inputs.shape)
      layer.kernel = spectral_normalize(layer.kernel)
      outputs = layer(inputs)
  
    Args:
      name_filter: Optionally, a method that takes a Variable name as input and
        returns whether this Variable should be normalized.
      power_iteration_rounds: The number of iterations of the power method to
        perform per step. A higher number yields a better approximation of the
        true spectral norm.
      equality_constrained: If set to `True` will normalize the matrix such that
        its spectral norm is equal to 1, otherwise, will normalize the matrix such
        that its norm is at most 1.
      training: Whether to update the spectral normalization on variable
        access. This is useful to turn off during eval, for example, to not affect
        the graph during evaluation.
  
    Returns:
      A custom getter function that applies Spectral Normalization to all
      Variables whose names match `name_filter`.
  
    Raises:
      ValueError: If name_filter is not callable.
    """
    if not callable(name_filter):
      raise ValueError('name_filter must be callable')
  
    def _internal_getter(getter, name, *args, **kwargs):
      """A custom getter function that applies Spectral Normalization.
  
      Args:
        getter: The true getter to call.
        name: Name of new/existing variable, in the same format as
          tf.get_variable.
        *args: Other positional arguments, in the same format as tf.get_variable.
        **kwargs: Keyword arguments, in the same format as tf.get_variable.
  
      Returns:
        The return value of `getter(name, *args, **kwargs)`, spectrally
        normalized.
  
      Raises:
        ValueError: If used incorrectly, or if `dtype` is not supported.
      """
      if not name_filter(name):
        return getter(name, *args, **kwargs)
  
      if name.endswith(_PERSISTED_U_VARIABLE_SUFFIX):
        raise ValueError(
            'Cannot apply Spectral Normalization to internal variables created '
            'for Spectral Normalization. Tried to normalized variable [%s]' %
            name)
  
      if kwargs['dtype'] not in _OK_DTYPES_FOR_SPECTRAL_NORM:
        raise ValueError('Disallowed data type {}'.format(kwargs['dtype']))
  
      # This layer's weight Variable/PartitionedVariable.
      w_tensor = getter(name, *args, **kwargs)
  
      if len(w_tensor.get_shape()) < 2:
        raise ValueError(
            'Spectral norm can only be applied to multi-dimensional tensors')
  
      return spectral_normalize(
          w_tensor,
          power_iteration_rounds=power_iteration_rounds,
          equality_constrained=equality_constrained,
          training=training,
          name=(name + '/spectral_normalize'))
  
    return _internal_getter
  