keras/src/layers/normalization/group_normalization.py/GroupNormalization/_apply_normalization
class GroupNormalization:    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
        )

        # Compute the batch normalization.
        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