group normalization layer group normalization divides the channels into groups and computes within each group the mean and variance for normalization empirically its accuracy is more stable than batch norm in a wide range of small batch sizes if learning rate is adjusted linearly with batch sizes relation to layer normalization if the number of groups is set to 1 then this operation becomes nearly identical to layer normalization see layer normalization docs for details relation to instance normalization if the number of groups is set to the input dimension number of groups is equal to number of channels then this operation becomes identical to instance normalization you can achieve this via groups1 args groups integer the number of groups for group normalization can be in the range 1 n where n is the input dimension the input dimension must be divisible by the number of groups defaults to 32 axis integer or listtuple the axis or axes to normalize across typically this is the features axisaxes the leftout axes are typically the batch axisaxes 1 is the last dimension in the input defaults to 1 epsilon small float added to variance to avoid dividing by zero defaults to 1e3 center if true add offset of beta to normalized tensor if false beta is ignored defaults to true scale if true multiply by gamma if false gamma is not used when the next layer is linear also e g relu this can be disabled since the scaling will be done by the next layer defaults to true betainitializer initializer for the beta weight defaults to zeros gammainitializer initializer for the gamma weight defaults to ones betaregularizer optional regularizer for the beta weight none by default gammaregularizer optional regularizer for the gamma weight none by default betaconstraint optional constraint for the beta weight none by default gammaconstraint optional constraint for the gamma weight none by default input shape arbitrary use the keyword argument inputshape tuple of integers does not include the samples axis when using this layer as the first layer in a model output shape same shape as input kwargs base layer keyword arguments e g name and dtype reference yuxin wu kaiming he 2018https arxiv orgabs1803 08494 compute the batch normalization group normalization layer group normalization divides the channels into groups and computes within each group the mean and variance for normalization empirically its accuracy is more stable than batch norm in a wide range of small batch sizes if learning rate is adjusted linearly with batch sizes relation to layer normalization if the number of groups is set to 1 then this operation becomes nearly identical to layer normalization see layer normalization docs for details relation to instance normalization if the number of groups is set to the input dimension number of groups is equal to number of channels then this operation becomes identical to instance normalization you can achieve this via groups 1 args groups integer the number of groups for group normalization can be in the range 1 n where n is the input dimension the input dimension must be divisible by the number of groups defaults to 32 axis integer or list tuple the axis or axes to normalize across typically this is the features axis axes the left out axes are typically the batch axis axes 1 is the last dimension in the input defaults to 1 epsilon small float added to variance to avoid dividing by zero defaults to 1e 3 center if true add offset of beta to normalized tensor if false beta is ignored defaults to true scale if true multiply by gamma if false gamma is not used when the next layer is linear also e g relu this can be disabled since the scaling will be done by the next layer defaults to true beta_initializer initializer for the beta weight defaults to zeros gamma_initializer initializer for the gamma weight defaults to ones beta_regularizer optional regularizer for the beta weight none by default gamma_regularizer optional regularizer for the gamma weight none by default beta_constraint optional constraint for the beta weight none by default gamma_constraint optional constraint for the gamma weight none by default input shape arbitrary use the keyword argument input_shape tuple of integers does not include the samples axis when using this layer as the first layer in a model output shape same shape as input kwargs base layer keyword arguments e g name and dtype reference yuxin wu kaiming he 2018 https arxiv org abs 1803 08494 compute the batch normalization
from keras import constraints
from keras import initializers
from keras import ops
from keras import regularizers
from keras.api_export import keras_export
from keras.layers.input_spec import InputSpec
from keras.layers.layer import Layer


@keras_export("keras.layers.GroupNormalization")
class GroupNormalization(Layer):
    

    def __init__(
        self,
        groups=32,
        axis=-1,
        epsilon=1e-3,
        center=True,
        scale=True,
        beta_initializer="zeros",
        gamma_initializer="ones",
        beta_regularizer=None,
        gamma_regularizer=None,
        beta_constraint=None,
        gamma_constraint=None,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.supports_masking = True
        self.groups = groups
        self.axis = axis
        self.epsilon = epsilon
        self.center = center
        self.scale = scale
        self.beta_initializer = initializers.get(beta_initializer)
        self.gamma_initializer = initializers.get(gamma_initializer)
        self.beta_regularizer = regularizers.get(beta_regularizer)
        self.gamma_regularizer = regularizers.get(gamma_regularizer)
        self.beta_constraint = constraints.get(beta_constraint)
        self.gamma_constraint = constraints.get(gamma_constraint)

    def build(self, input_shape):
        dim = input_shape[self.axis]

        if dim is None:
            raise ValueError(
                f"Axis {self.axis} of input tensor should have a defined "
                "dimension but the layer received an input with shape "
                f"{input_shape}."
            )

        if self.groups == -1:
            self.groups = dim

        if dim < self.groups:
            raise ValueError(
                f"Number of groups ({self.groups}) cannot be more than the "
                f"number of channels ({dim})."
            )

        if dim % self.groups != 0:
            raise ValueError(
                f"Number of groups ({self.groups}) must be a multiple "
                f"of the number of channels ({dim})."
            )

        self.input_spec = InputSpec(
            ndim=len(input_shape), axes={self.axis: dim}
        )

        if self.scale:
            self.gamma = self.add_weight(
                shape=(dim,),
                name="gamma",
                initializer=self.gamma_initializer,
                regularizer=self.gamma_regularizer,
                constraint=self.gamma_constraint,
            )
        else:
            self.gamma = None

        if self.center:
            self.beta = self.add_weight(
                shape=(dim,),
                name="beta",
                initializer=self.beta_initializer,
                regularizer=self.beta_regularizer,
                constraint=self.beta_constraint,
            )
        else:
            self.beta = None

        super().build(input_shape)

    def call(self, inputs):
        reshaped_inputs = self._reshape_into_groups(inputs)
        normalized_inputs = self._apply_normalization(
            reshaped_inputs, inputs.shape
        )
        return ops.reshape(normalized_inputs, ops.shape(inputs))

    def _reshape_into_groups(self, inputs):
        input_shape = ops.shape(inputs)
        group_shape = list(inputs.shape)
        group_shape[0] = -1
        for i, e in enumerate(group_shape[1:]):
            if e is None:
                group_shape[i + 1] = input_shape[i + 1]

        group_shape[self.axis] = input_shape[self.axis] // self.groups
        group_shape.insert(self.axis, self.groups)
        reshaped_inputs = ops.reshape(inputs, group_shape)
        return reshaped_inputs

    def _apply_normalization(self, reshaped_inputs, input_shape):
        group_reduction_axes = list(range(1, len(reshaped_inputs.shape)))

        axis = -2 if self.axis == -1 else self.axis - 1
        group_reduction_axes.pop(axis)

        broadcast_shape = self._create_broadcast_shape(input_shape)
        mean, variance = ops.moments(
            reshaped_inputs, axes=group_reduction_axes, keepdims=True
        )

        
        inv = ops.rsqrt(variance + self.epsilon)
        if self.scale:
            gamma = ops.reshape(self.gamma, broadcast_shape)
            gamma = ops.cast(gamma, reshaped_inputs.dtype)
            inv = inv * gamma

        res = -mean * inv
        if self.center:
            beta = ops.reshape(self.beta, broadcast_shape)
            beta = ops.cast(beta, reshaped_inputs.dtype)
            res = res + beta

        normalized_inputs = reshaped_inputs * inv + res
        return normalized_inputs

    def _create_broadcast_shape(self, input_shape):
        broadcast_shape = [1] * len(input_shape)
        broadcast_shape[self.axis] = input_shape[self.axis] // self.groups
        broadcast_shape.insert(self.axis, self.groups)
        return broadcast_shape

    def compute_output_shape(self, input_shape):
        return input_shape

    def get_config(self):
        config = {
            "groups": self.groups,
            "axis": self.axis,
            "epsilon": self.epsilon,
            "center": self.center,
            "scale": self.scale,
            "beta_initializer": initializers.serialize(self.beta_initializer),
            "gamma_initializer": initializers.serialize(self.gamma_initializer),
            "beta_regularizer": regularizers.serialize(self.beta_regularizer),
            "gamma_regularizer": regularizers.serialize(self.gamma_regularizer),
            "beta_constraint": constraints.serialize(self.beta_constraint),
            "gamma_constraint": constraints.serialize(self.gamma_constraint),
        }
        base_config = super().get_config()
        return {**base_config, **config}