performs spectral normalization on the weights of a target layer this wrapper controls the lipschitz constant of the weights of a layer by constraining their spectral norm which can stabilize the training of gans args layer a keras layers layer instance that has either a kernel e g conv2d dense or an embeddings attribute embedding layer poweriterations int the number of iterations during normalization kwargs base wrapper keyword arguments examples wrap keras layers conv2d x np random rand1 10 10 1 conv2d spectralnormalizationkeras layers conv2d2 2 y conv2dx y shape 1 9 9 2 wrap keras layers dense x np random rand1 10 10 1 dense spectralnormalizationkeras layers dense10 y densex y shape 1 10 10 10 reference spectral normalization for ganhttps arxiv orgabs1802 05957 generate spectral normalized weights this method will update the value of self kernel with the spectral normalized value so that the layer is ready for call check for zeroes weights vectoru tf stopgradientvectoru vectorv tf stopgradientvectorv performs spectral normalization on the weights of a target layer this wrapper controls the lipschitz constant of the weights of a layer by constraining their spectral norm which can stabilize the training of gans args layer a keras layers layer instance that has either a kernel e g conv2d dense or an embeddings attribute embedding layer power_iterations int the number of iterations during normalization kwargs base wrapper keyword arguments examples wrap keras layers conv2d x np random rand 1 10 10 1 conv2d spectralnormalization keras layers conv2d 2 2 y conv2d x y shape 1 9 9 2 wrap keras layers dense x np random rand 1 10 10 1 dense spectralnormalization keras layers dense 10 y dense x y shape 1 10 10 10 reference spectral normalization for gan https arxiv org abs 1802 05957 generate spectral normalized weights this method will update the value of self kernel with the spectral normalized value so that the layer is ready for call check for zeroes weights vector_u tf stop_gradient vector_u vector_v tf stop_gradient vector_v
from keras import initializers
from keras import ops
from keras.api_export import keras_export
from keras.layers import Wrapper
from keras.layers.input_spec import InputSpec
from keras.utils.numerical_utils import normalize


@keras_export("keras.layers.SpectralNormalization")
class SpectralNormalization(Wrapper):
    

    def __init__(self, layer, power_iterations=1, **kwargs):
        super().__init__(layer, **kwargs)
        if power_iterations <= 0:
            raise ValueError(
                "`power_iterations` should be greater than zero. Received: "
                f"`power_iterations={power_iterations}`"
            )
        self.power_iterations = power_iterations

    def build(self, input_shape):
        super().build(input_shape)
        self.input_spec = InputSpec(shape=[None] + list(input_shape[1:]))

        if hasattr(self.layer, "kernel"):
            self.kernel = self.layer.kernel
        elif hasattr(self.layer, "embeddings"):
            self.kernel = self.layer.embeddings
        else:
            raise ValueError(
                f"{type(self.layer).__name__} object has no attribute 'kernel' "
                "nor 'embeddings'"
            )

        self.kernel_shape = self.kernel.shape

        self.vector_u = self.add_weight(
            shape=(1, self.kernel_shape[-1]),
            initializer=initializers.TruncatedNormal(stddev=0.02),
            trainable=False,
            name="vector_u",
            dtype=self.kernel.dtype,
        )

    def call(self, inputs, training=False):
        if training:
            self.normalize_weights()

        output = self.layer(inputs)
        return ops.cast(output, inputs.dtype)

    def compute_output_shape(self, input_shape):
        return self.layer.compute_output_shape(input_shape)

    def normalize_weights(self):
        

        weights = ops.reshape(self.kernel, [-1, self.kernel_shape[-1]])
        vector_u = self.vector_u

        
        if not all([w == 0.0 for w in weights]):
            for _ in range(self.power_iterations):
                vector_v = normalize(
                    ops.matmul(vector_u, ops.transpose(weights)), axis=None
                )
                vector_u = normalize(ops.matmul(vector_v, weights), axis=None)
            
            
            sigma = ops.matmul(
                ops.matmul(vector_v, weights), ops.transpose(vector_u)
            )
            self.vector_u.assign(ops.cast(vector_u, self.vector_u.dtype))
            self.kernel.assign(
                ops.cast(
                    ops.reshape(self.kernel / sigma, self.kernel_shape),
                    self.kernel.dtype,
                )
            )

    def get_config(self):
        config = {"power_iterations": self.power_iterations}
        base_config = super().get_config()
        return {**base_config, **config}