We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Layernorm not supporting axis [-2, 3]
Hi, 
I wanted to normalise my output on the -2 and -3 axis, (image height and width), however, it seems that the with rms_scaling=true, the self.gamma is not broadcasted to same shape as layer input causing this error, 
```
inputs shape: (1, 1920, 1200, 3)
inv shape: (1, 1, 1, 3)
gamma_cast shape: (1920, 1200)
inv shape: (1, 1920, 1200, 3)
2024-04-30 13:50:54.238379: W tensorflow/core/framework/local_rendezvous.cc:404] Local rendezvous is aborting with status: INVALID_ARGUMENT: Incompatible shapes: [1,1920,1200,3] vs. [1920,1200]
Traceback (most recent call last):
  File "C:\Users\88bbh\PycharmProjects\AI\tempt.py", line 10, in <module>
    layer(np.zeros((1, 1920, 1200, 3)))
  File "C:\Users\88bbh\PycharmProjects\AI\venv\lib\site-packages\keras\src\utils\traceback_utils.py", line 122, in error_handler
    raise e.with_traceback(filtered_tb) from None
  File "C:\Users\88bbh\PycharmProjects\AI\venv\lib\site-packages\tensorflow\python\framework\ops.py", line 5983, in raise_from_not_ok_status
    raise core._status_to_exception(e) from None  # pylint: disable=protected-access
tensorflow.python.framework.errors_impl.InvalidArgumentError: Exception encountered when calling LayerNormalization.call().

{{function_node __wrapped__Mul_device_/job:localhost/replica:0/task:0/device:CPU:0}} Incompatible shapes: [1,1920,1200,3] vs. [1920,1200] [Op:Mul] name: 

Arguments received by LayerNormalization.call():
  • inputs=tf.Tensor(shape=(1, 1920, 1200, 3), dtype=float32)
```

code to reproduce
```
layer = keras.layers.LayerNormalization(axis=[-3, -2], rms_scaling=True)
layer.build([None, 1920, 1200, 3])
layer(np.zeros((1, 1920, 1200, 3)))
```

the error is in layernorm call method
```
        if self.rms_scaling:
            # Calculate outputs with only variance and gamma if rms scaling
            # is enabled
            # Calculate the variance along self.axis (layer activations).
            variance = ops.var(inputs, axis=self.axis, keepdims=True)
            inv = ops.rsqrt(variance + self.epsilon)
            print("inputs shape:", inputs.shape)
            print("inv shape:", inv.shape)
            print("gamma_cast shape:", self.gamma.shape)
            print("inv shape:", (inputs * inv).shape)
            outputs = inputs * inv * ops.cast(self.gamma, inputs.dtype)
```

the error can be fixed by changing
```
outputs = inputs * inv * ops.cast(self.gamma, inputs.dtype)
to
outputs = inputs * inv * ops.cast(_broadcast(self.gamma), inputs.dtype)
```
please fix it in the next update
thank you

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `keras/src/layers/normalization/layer_normalization.py`
Content:
```
1 from keras.src import constraints
2 from keras.src import initializers
3 from keras.src import ops
4 from keras.src import regularizers
5 from keras.src.api_export import keras_export
6 from keras.src.layers.layer import Layer
7 
8 
9 @keras_export("keras.layers.LayerNormalization")
10 class LayerNormalization(Layer):
11     """Layer normalization layer (Ba et al., 2016).
12 
13     Normalize the activations of the previous layer for each given example in a
14     batch independently, rather than across a batch like Batch Normalization.
15     i.e. applies a transformation that maintains the mean activation within each
16     example close to 0 and the activation standard deviation close to 1.
17 
18     If `scale` or `center` are enabled, the layer will scale the normalized
19     outputs by broadcasting them with a trainable variable `gamma`, and center
20     the outputs by broadcasting with a trainable variable `beta`. `gamma` will
21     default to a ones tensor and `beta` will default to a zeros tensor, so that
22     centering and scaling are no-ops before training has begun.
23 
24     So, with scaling and centering enabled the normalization equations
25     are as follows:
26 
27     Let the intermediate activations for a mini-batch to be the `inputs`.
28 
29     For each sample `x_i` in `inputs` with `k` features, we compute the mean and
30     variance of the sample:
31 
32     ```python
33     mean_i = sum(x_i[j] for j in range(k)) / k
34     var_i = sum((x_i[j] - mean_i) ** 2 for j in range(k)) / k
35     ```
36 
37     and then compute a normalized `x_i_normalized`, including a small factor
38     `epsilon` for numerical stability.
39 
40     ```python
41     x_i_normalized = (x_i - mean_i) / sqrt(var_i + epsilon)
42     ```
43 
44     And finally `x_i_normalized ` is linearly transformed by `gamma` and `beta`,
45     which are learned parameters:
46 
47     ```python
48     output_i = x_i_normalized * gamma + beta
49     ```
50 
51     `gamma` and `beta` will span the axes of `inputs` specified in `axis`, and
52     this part of the inputs' shape must be fully defined.
53 
54     For example:
55 
56     >>> layer = keras.layers.LayerNormalization(axis=[1, 2, 3])
57     >>> layer.build([5, 20, 30, 40])
58     >>> print(layer.beta.shape)
59     (20, 30, 40)
60     >>> print(layer.gamma.shape)
61     (20, 30, 40)
62 
63     Note that other implementations of layer normalization may choose to define
64     `gamma` and `beta` over a separate set of axes from the axes being
65     normalized across. For example, Group Normalization
66     ([Wu et al. 2018](https://arxiv.org/abs/1803.08494)) with group size of 1
67     corresponds to a Layer Normalization that normalizes across height, width,
68     and channel and has `gamma` and `beta` span only the channel dimension.
69     So, this Layer Normalization implementation will not match a Group
70     Normalization layer with group size set to 1.
71 
72     Args:
73         axis: Integer or List/Tuple. The axis or axes to normalize across.
74             Typically, this is the features axis/axes. The left-out axes are
75             typically the batch axis/axes. `-1` is the last dimension in the
76             input. Defaults to `-1`.
77         epsilon: Small float added to variance to avoid dividing by zero.
78             Defaults to 1e-3.
79         center: If True, add offset of `beta` to normalized tensor. If False,
80             `beta` is ignored. Defaults to `True`.
81         scale: If True, multiply by `gamma`. If False, `gamma` is not used.
82             When the next layer is linear (also e.g. `nn.relu`), this can be
83             disabled since the scaling will be done by the next layer.
84             Defaults to `True`.
85         rms_scaling: If True, `center` and `scale` are ignored, and the
86             inputs are scaled by `gamma` and the inverse square root
87             of the square of all inputs. This is an approximate and faster
88             approach that avoids ever computing the mean of the input.
89         beta_initializer: Initializer for the beta weight. Defaults to zeros.
90         gamma_initializer: Initializer for the gamma weight. Defaults to ones.
91         beta_regularizer: Optional regularizer for the beta weight.
92             None by default.
93         gamma_regularizer: Optional regularizer for the gamma weight.
94             None by default.
95         beta_constraint: Optional constraint for the beta weight.
96             None by default.
97         gamma_constraint: Optional constraint for the gamma weight.
98             None by default.
99         **kwargs: Base layer keyword arguments (e.g. `name` and `dtype`).
100 
101 
102     Reference:
103 
104     - [Lei Ba et al., 2016](https://arxiv.org/abs/1607.06450).
105     """
106 
107     def __init__(
108         self,
109         axis=-1,
110         epsilon=1e-3,
111         center=True,
112         scale=True,
113         rms_scaling=False,
114         beta_initializer="zeros",
115         gamma_initializer="ones",
116         beta_regularizer=None,
117         gamma_regularizer=None,
118         beta_constraint=None,
119         gamma_constraint=None,
120         **kwargs
121     ):
122         super().__init__(**kwargs)
123         if isinstance(axis, (list, tuple)):
124             self.axis = list(axis)
125         elif isinstance(axis, int):
126             self.axis = axis
127         else:
128             raise TypeError(
129                 "Expected an int or a list/tuple of ints for the "
130                 "argument 'axis', but received: %r" % axis
131             )
132 
133         self.epsilon = epsilon
134         self.center = center
135         self.scale = scale
136         self.rms_scaling = rms_scaling
137         self.beta_initializer = initializers.get(beta_initializer)
138         self.gamma_initializer = initializers.get(gamma_initializer)
139         self.beta_regularizer = regularizers.get(beta_regularizer)
140         self.gamma_regularizer = regularizers.get(gamma_regularizer)
141         self.beta_constraint = constraints.get(beta_constraint)
142         self.gamma_constraint = constraints.get(gamma_constraint)
143 
144         self.supports_masking = True
145         self.autocast = False
146 
147     def build(self, input_shape):
148         if isinstance(self.axis, list):
149             shape = tuple([input_shape[dim] for dim in self.axis])
150         else:
151             shape = (input_shape[self.axis],)
152             self.axis = [self.axis]
153         if self.scale or self.rms_scaling:
154             self.gamma = self.add_weight(
155                 name="gamma",
156                 shape=shape,
157                 initializer=self.gamma_initializer,
158                 regularizer=self.gamma_regularizer,
159                 constraint=self.gamma_constraint,
160                 trainable=True,
161                 autocast=False,
162             )
163         else:
164             self.gamma = None
165 
166         if self.center and not self.rms_scaling:
167             self.beta = self.add_weight(
168                 name="beta",
169                 shape=shape,
170                 initializer=self.beta_initializer,
171                 regularizer=self.beta_regularizer,
172                 constraint=self.beta_constraint,
173                 trainable=True,
174                 autocast=False,
175             )
176         else:
177             self.beta = None
178 
179         self.built = True
180 
181     def call(self, inputs):
182         inputs = ops.cast(inputs, self.compute_dtype)
183         # Compute the axes along which to reduce the mean / variance
184         input_shape = inputs.shape
185         ndims = len(input_shape)
186 
187         # Broadcasting only necessary for norm when the axis is not just
188         # the last dimension
189         broadcast_shape = [1] * ndims
190         for dim in self.axis:
191             broadcast_shape[dim] = input_shape[dim]
192 
193         def _broadcast(v):
194             if (
195                 v is not None
196                 and len(v.shape) != ndims
197                 and self.axis != [ndims - 1]
198             ):
199                 return ops.reshape(v, broadcast_shape)
200             return v
201 
202         input_dtype = inputs.dtype
203         if input_dtype in ("float16", "bfloat16") and self.dtype == "float32":
204             # If mixed precision is used, cast inputs to float32 so that
205             # this is at least as numerically stable as the fused version.
206             inputs = ops.cast(inputs, "float32")
207 
208         if self.rms_scaling:
209             # Calculate outputs with only variance and gamma if rms scaling
210             # is enabled
211             # Calculate the variance along self.axis (layer activations).
212             variance = ops.var(inputs, axis=self.axis, keepdims=True)
213             inv = ops.rsqrt(variance + self.epsilon)
214 
215             outputs = inputs * inv * ops.cast(self.gamma, inputs.dtype)
216         else:
217             # Calculate the mean & variance along self.axis (layer activations).
218             mean, variance = ops.moments(inputs, axes=self.axis, keepdims=True)
219             gamma, beta = _broadcast(self.gamma), _broadcast(self.beta)
220 
221             inv = ops.rsqrt(variance + self.epsilon)
222             if gamma is not None:
223                 gamma = ops.cast(gamma, inputs.dtype)
224                 inv = inv * gamma
225 
226             res = -mean * inv
227             if beta is not None:
228                 beta = ops.cast(beta, inputs.dtype)
229                 res = res + beta
230 
231             outputs = inputs * inv + res
232 
233         return ops.cast(outputs, input_dtype)
234 
235     def compute_output_shape(self, input_shape):
236         return input_shape
237 
238     def get_config(self):
239         config = {
240             "axis": self.axis,
241             "epsilon": self.epsilon,
242             "center": self.center,
243             "scale": self.scale,
244             "beta_initializer": initializers.serialize(self.beta_initializer),
245             "gamma_initializer": initializers.serialize(self.gamma_initializer),
246             "beta_regularizer": regularizers.serialize(self.beta_regularizer),
247             "gamma_regularizer": regularizers.serialize(self.gamma_regularizer),
248             "beta_constraint": constraints.serialize(self.beta_constraint),
249             "gamma_constraint": constraints.serialize(self.gamma_constraint),
250         }
251         base_config = super().get_config()
252         return {**base_config, **config}
253
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/keras/src/layers/normalization/layer_normalization.py b/keras/src/layers/normalization/layer_normalization.py
--- a/keras/src/layers/normalization/layer_normalization.py
+++ b/keras/src/layers/normalization/layer_normalization.py
@@ -212,7 +212,9 @@
             variance = ops.var(inputs, axis=self.axis, keepdims=True)
             inv = ops.rsqrt(variance + self.epsilon)
 
-            outputs = inputs * inv * ops.cast(self.gamma, inputs.dtype)
+            outputs = (
+                inputs * inv * ops.cast(_broadcast(self.gamma), inputs.dtype)
+            )
         else:
             # Calculate the mean & variance along self.axis (layer activations).
             mean, variance = ops.moments(inputs, axes=self.axis, keepdims=True)

{"golden_diff": "diff --git a/keras/src/layers/normalization/layer_normalization.py b/keras/src/layers/normalization/layer_normalization.py\n--- a/keras/src/layers/normalization/layer_normalization.py\n+++ b/keras/src/layers/normalization/layer_normalization.py\n@@ -212,7 +212,9 @@\n             variance = ops.var(inputs, axis=self.axis, keepdims=True)\n             inv = ops.rsqrt(variance + self.epsilon)\n \n-            outputs = inputs * inv * ops.cast(self.gamma, inputs.dtype)\n+            outputs = (\n+                inputs * inv * ops.cast(_broadcast(self.gamma), inputs.dtype)\n+            )\n         else:\n             # Calculate the mean & variance along self.axis (layer activations).\n             mean, variance = ops.moments(inputs, axes=self.axis, keepdims=True)\n", "issue": "Layernorm not supporting axis [-2, 3]\nHi, \r\nI wanted to normalise my output on the -2 and -3 axis, (image height and width), however, it seems that the with rms_scaling=true, the self.gamma is not broadcasted to same shape as layer input causing this error, \r\n```\r\ninputs shape: (1, 1920, 1200, 3)\r\ninv shape: (1, 1, 1, 3)\r\ngamma_cast shape: (1920, 1200)\r\ninv shape: (1, 1920, 1200, 3)\r\n2024-04-30 13:50:54.238379: W tensorflow/core/framework/local_rendezvous.cc:404] Local rendezvous is aborting with status: INVALID_ARGUMENT: Incompatible shapes: [1,1920,1200,3] vs. [1920,1200]\r\nTraceback (most recent call last):\r\n  File \"C:\\Users\\88bbh\\PycharmProjects\\AI\\tempt.py\", line 10, in <module>\r\n    layer(np.zeros((1, 1920, 1200, 3)))\r\n  File \"C:\\Users\\88bbh\\PycharmProjects\\AI\\venv\\lib\\site-packages\\keras\\src\\utils\\traceback_utils.py\", line 122, in error_handler\r\n    raise e.with_traceback(filtered_tb) from None\r\n  File \"C:\\Users\\88bbh\\PycharmProjects\\AI\\venv\\lib\\site-packages\\tensorflow\\python\\framework\\ops.py\", line 5983, in raise_from_not_ok_status\r\n    raise core._status_to_exception(e) from None  # pylint: disable=protected-access\r\ntensorflow.python.framework.errors_impl.InvalidArgumentError: Exception encountered when calling LayerNormalization.call().\r\n\r\n{{function_node __wrapped__Mul_device_/job:localhost/replica:0/task:0/device:CPU:0}} Incompatible shapes: [1,1920,1200,3] vs. [1920,1200] [Op:Mul] name: \r\n\r\nArguments received by LayerNormalization.call():\r\n  \u2022 inputs=tf.Tensor(shape=(1, 1920, 1200, 3), dtype=float32)\r\n```\r\n\r\ncode to reproduce\r\n```\r\nlayer = keras.layers.LayerNormalization(axis=[-3, -2], rms_scaling=True)\r\nlayer.build([None, 1920, 1200, 3])\r\nlayer(np.zeros((1, 1920, 1200, 3)))\r\n```\r\n\r\nthe error is in layernorm call method\r\n```\r\n        if self.rms_scaling:\r\n            # Calculate outputs with only variance and gamma if rms scaling\r\n            # is enabled\r\n            # Calculate the variance along self.axis (layer activations).\r\n            variance = ops.var(inputs, axis=self.axis, keepdims=True)\r\n            inv = ops.rsqrt(variance + self.epsilon)\r\n            print(\"inputs shape:\", inputs.shape)\r\n            print(\"inv shape:\", inv.shape)\r\n            print(\"gamma_cast shape:\", self.gamma.shape)\r\n            print(\"inv shape:\", (inputs * inv).shape)\r\n            outputs = inputs * inv * ops.cast(self.gamma, inputs.dtype)\r\n```\r\n\r\nthe error can be fixed by changing\r\n```\r\noutputs = inputs * inv * ops.cast(self.gamma, inputs.dtype)\r\nto\r\noutputs = inputs * inv * ops.cast(_broadcast(self.gamma), inputs.dtype)\r\n```\r\nplease fix it in the next update\r\nthank you\n", "before_files": [{"content": "from keras.src import constraints\nfrom keras.src import initializers\nfrom keras.src import ops\nfrom keras.src import regularizers\nfrom keras.src.api_export import keras_export\nfrom keras.src.layers.layer import Layer\n\n\n@keras_export(\"keras.layers.LayerNormalization\")\nclass LayerNormalization(Layer):\n    \"\"\"Layer normalization layer (Ba et al., 2016).\n\n    Normalize the activations of the previous layer for each given example in a\n    batch independently, rather than across a batch like Batch Normalization.\n    i.e. applies a transformation that maintains the mean activation within each\n    example close to 0 and the activation standard deviation close to 1.\n\n    If `scale` or `center` are enabled, the layer will scale the normalized\n    outputs by broadcasting them with a trainable variable `gamma`, and center\n    the outputs by broadcasting with a trainable variable `beta`. `gamma` will\n    default to a ones tensor and `beta` will default to a zeros tensor, so that\n    centering and scaling are no-ops before training has begun.\n\n    So, with scaling and centering enabled the normalization equations\n    are as follows:\n\n    Let the intermediate activations for a mini-batch to be the `inputs`.\n\n    For each sample `x_i` in `inputs` with `k` features, we compute the mean and\n    variance of the sample:\n\n    ```python\n    mean_i = sum(x_i[j] for j in range(k)) / k\n    var_i = sum((x_i[j] - mean_i) ** 2 for j in range(k)) / k\n    ```\n\n    and then compute a normalized `x_i_normalized`, including a small factor\n    `epsilon` for numerical stability.\n\n    ```python\n    x_i_normalized = (x_i - mean_i) / sqrt(var_i + epsilon)\n    ```\n\n    And finally `x_i_normalized ` is linearly transformed by `gamma` and `beta`,\n    which are learned parameters:\n\n    ```python\n    output_i = x_i_normalized * gamma + beta\n    ```\n\n    `gamma` and `beta` will span the axes of `inputs` specified in `axis`, and\n    this part of the inputs' shape must be fully defined.\n\n    For example:\n\n    >>> layer = keras.layers.LayerNormalization(axis=[1, 2, 3])\n    >>> layer.build([5, 20, 30, 40])\n    >>> print(layer.beta.shape)\n    (20, 30, 40)\n    >>> print(layer.gamma.shape)\n    (20, 30, 40)\n\n    Note that other implementations of layer normalization may choose to define\n    `gamma` and `beta` over a separate set of axes from the axes being\n    normalized across. For example, Group Normalization\n    ([Wu et al. 2018](https://arxiv.org/abs/1803.08494)) with group size of 1\n    corresponds to a Layer Normalization that normalizes across height, width,\n    and channel and has `gamma` and `beta` span only the channel dimension.\n    So, this Layer Normalization implementation will not match a Group\n    Normalization layer with group size set to 1.\n\n    Args:\n        axis: Integer or List/Tuple. The axis or axes to normalize across.\n            Typically, this is the features axis/axes. The left-out axes are\n            typically the batch axis/axes. `-1` is the last dimension in the\n            input. Defaults to `-1`.\n        epsilon: Small float added to variance to avoid dividing by zero.\n            Defaults to 1e-3.\n        center: If True, add offset of `beta` to normalized tensor. If False,\n            `beta` is ignored. Defaults to `True`.\n        scale: If True, multiply by `gamma`. If False, `gamma` is not used.\n            When the next layer is linear (also e.g. `nn.relu`), this can be\n            disabled since the scaling will be done by the next layer.\n            Defaults to `True`.\n        rms_scaling: If True, `center` and `scale` are ignored, and the\n            inputs are scaled by `gamma` and the inverse square root\n            of the square of all inputs. This is an approximate and faster\n            approach that avoids ever computing the mean of the input.\n        beta_initializer: Initializer for the beta weight. Defaults to zeros.\n        gamma_initializer: Initializer for the gamma weight. Defaults to ones.\n        beta_regularizer: Optional regularizer for the beta weight.\n            None by default.\n        gamma_regularizer: Optional regularizer for the gamma weight.\n            None by default.\n        beta_constraint: Optional constraint for the beta weight.\n            None by default.\n        gamma_constraint: Optional constraint for the gamma weight.\n            None by default.\n        **kwargs: Base layer keyword arguments (e.g. `name` and `dtype`).\n\n\n    Reference:\n\n    - [Lei Ba et al., 2016](https://arxiv.org/abs/1607.06450).\n    \"\"\"\n\n    def __init__(\n        self,\n        axis=-1,\n        epsilon=1e-3,\n        center=True,\n        scale=True,\n        rms_scaling=False,\n        beta_initializer=\"zeros\",\n        gamma_initializer=\"ones\",\n        beta_regularizer=None,\n        gamma_regularizer=None,\n        beta_constraint=None,\n        gamma_constraint=None,\n        **kwargs\n    ):\n        super().__init__(**kwargs)\n        if isinstance(axis, (list, tuple)):\n            self.axis = list(axis)\n        elif isinstance(axis, int):\n            self.axis = axis\n        else:\n            raise TypeError(\n                \"Expected an int or a list/tuple of ints for the \"\n                \"argument 'axis', but received: %r\" % axis\n            )\n\n        self.epsilon = epsilon\n        self.center = center\n        self.scale = scale\n        self.rms_scaling = rms_scaling\n        self.beta_initializer = initializers.get(beta_initializer)\n        self.gamma_initializer = initializers.get(gamma_initializer)\n        self.beta_regularizer = regularizers.get(beta_regularizer)\n        self.gamma_regularizer = regularizers.get(gamma_regularizer)\n        self.beta_constraint = constraints.get(beta_constraint)\n        self.gamma_constraint = constraints.get(gamma_constraint)\n\n        self.supports_masking = True\n        self.autocast = False\n\n    def build(self, input_shape):\n        if isinstance(self.axis, list):\n            shape = tuple([input_shape[dim] for dim in self.axis])\n        else:\n            shape = (input_shape[self.axis],)\n            self.axis = [self.axis]\n        if self.scale or self.rms_scaling:\n            self.gamma = self.add_weight(\n                name=\"gamma\",\n                shape=shape,\n                initializer=self.gamma_initializer,\n                regularizer=self.gamma_regularizer,\n                constraint=self.gamma_constraint,\n                trainable=True,\n                autocast=False,\n            )\n        else:\n            self.gamma = None\n\n        if self.center and not self.rms_scaling:\n            self.beta = self.add_weight(\n                name=\"beta\",\n                shape=shape,\n                initializer=self.beta_initializer,\n                regularizer=self.beta_regularizer,\n                constraint=self.beta_constraint,\n                trainable=True,\n                autocast=False,\n            )\n        else:\n            self.beta = None\n\n        self.built = True\n\n    def call(self, inputs):\n        inputs = ops.cast(inputs, self.compute_dtype)\n        # Compute the axes along which to reduce the mean / variance\n        input_shape = inputs.shape\n        ndims = len(input_shape)\n\n        # Broadcasting only necessary for norm when the axis is not just\n        # the last dimension\n        broadcast_shape = [1] * ndims\n        for dim in self.axis:\n            broadcast_shape[dim] = input_shape[dim]\n\n        def _broadcast(v):\n            if (\n                v is not None\n                and len(v.shape) != ndims\n                and self.axis != [ndims - 1]\n            ):\n                return ops.reshape(v, broadcast_shape)\n            return v\n\n        input_dtype = inputs.dtype\n        if input_dtype in (\"float16\", \"bfloat16\") and self.dtype == \"float32\":\n            # If mixed precision is used, cast inputs to float32 so that\n            # this is at least as numerically stable as the fused version.\n            inputs = ops.cast(inputs, \"float32\")\n\n        if self.rms_scaling:\n            # Calculate outputs with only variance and gamma if rms scaling\n            # is enabled\n            # Calculate the variance along self.axis (layer activations).\n            variance = ops.var(inputs, axis=self.axis, keepdims=True)\n            inv = ops.rsqrt(variance + self.epsilon)\n\n            outputs = inputs * inv * ops.cast(self.gamma, inputs.dtype)\n        else:\n            # Calculate the mean & variance along self.axis (layer activations).\n            mean, variance = ops.moments(inputs, axes=self.axis, keepdims=True)\n            gamma, beta = _broadcast(self.gamma), _broadcast(self.beta)\n\n            inv = ops.rsqrt(variance + self.epsilon)\n            if gamma is not None:\n                gamma = ops.cast(gamma, inputs.dtype)\n                inv = inv * gamma\n\n            res = -mean * inv\n            if beta is not None:\n                beta = ops.cast(beta, inputs.dtype)\n                res = res + beta\n\n            outputs = inputs * inv + res\n\n        return ops.cast(outputs, input_dtype)\n\n    def compute_output_shape(self, input_shape):\n        return input_shape\n\n    def get_config(self):\n        config = {\n            \"axis\": self.axis,\n            \"epsilon\": self.epsilon,\n            \"center\": self.center,\n            \"scale\": self.scale,\n            \"beta_initializer\": initializers.serialize(self.beta_initializer),\n            \"gamma_initializer\": initializers.serialize(self.gamma_initializer),\n            \"beta_regularizer\": regularizers.serialize(self.beta_regularizer),\n            \"gamma_regularizer\": regularizers.serialize(self.gamma_regularizer),\n            \"beta_constraint\": constraints.serialize(self.beta_constraint),\n            \"gamma_constraint\": constraints.serialize(self.gamma_constraint),\n        }\n        base_config = super().get_config()\n        return {**base_config, **config}\n", "path": "keras/src/layers/normalization/layer_normalization.py"}], "after_files": [{"content": "from keras.src import constraints\nfrom keras.src import initializers\nfrom keras.src import ops\nfrom keras.src import regularizers\nfrom keras.src.api_export import keras_export\nfrom keras.src.layers.layer import Layer\n\n\n@keras_export(\"keras.layers.LayerNormalization\")\nclass LayerNormalization(Layer):\n    \"\"\"Layer normalization layer (Ba et al., 2016).\n\n    Normalize the activations of the previous layer for each given example in a\n    batch independently, rather than across a batch like Batch Normalization.\n    i.e. applies a transformation that maintains the mean activation within each\n    example close to 0 and the activation standard deviation close to 1.\n\n    If `scale` or `center` are enabled, the layer will scale the normalized\n    outputs by broadcasting them with a trainable variable `gamma`, and center\n    the outputs by broadcasting with a trainable variable `beta`. `gamma` will\n    default to a ones tensor and `beta` will default to a zeros tensor, so that\n    centering and scaling are no-ops before training has begun.\n\n    So, with scaling and centering enabled the normalization equations\n    are as follows:\n\n    Let the intermediate activations for a mini-batch to be the `inputs`.\n\n    For each sample `x_i` in `inputs` with `k` features, we compute the mean and\n    variance of the sample:\n\n    ```python\n    mean_i = sum(x_i[j] for j in range(k)) / k\n    var_i = sum((x_i[j] - mean_i) ** 2 for j in range(k)) / k\n    ```\n\n    and then compute a normalized `x_i_normalized`, including a small factor\n    `epsilon` for numerical stability.\n\n    ```python\n    x_i_normalized = (x_i - mean_i) / sqrt(var_i + epsilon)\n    ```\n\n    And finally `x_i_normalized ` is linearly transformed by `gamma` and `beta`,\n    which are learned parameters:\n\n    ```python\n    output_i = x_i_normalized * gamma + beta\n    ```\n\n    `gamma` and `beta` will span the axes of `inputs` specified in `axis`, and\n    this part of the inputs' shape must be fully defined.\n\n    For example:\n\n    >>> layer = keras.layers.LayerNormalization(axis=[1, 2, 3])\n    >>> layer.build([5, 20, 30, 40])\n    >>> print(layer.beta.shape)\n    (20, 30, 40)\n    >>> print(layer.gamma.shape)\n    (20, 30, 40)\n\n    Note that other implementations of layer normalization may choose to define\n    `gamma` and `beta` over a separate set of axes from the axes being\n    normalized across. For example, Group Normalization\n    ([Wu et al. 2018](https://arxiv.org/abs/1803.08494)) with group size of 1\n    corresponds to a Layer Normalization that normalizes across height, width,\n    and channel and has `gamma` and `beta` span only the channel dimension.\n    So, this Layer Normalization implementation will not match a Group\n    Normalization layer with group size set to 1.\n\n    Args:\n        axis: Integer or List/Tuple. The axis or axes to normalize across.\n            Typically, this is the features axis/axes. The left-out axes are\n            typically the batch axis/axes. `-1` is the last dimension in the\n            input. Defaults to `-1`.\n        epsilon: Small float added to variance to avoid dividing by zero.\n            Defaults to 1e-3.\n        center: If True, add offset of `beta` to normalized tensor. If False,\n            `beta` is ignored. Defaults to `True`.\n        scale: If True, multiply by `gamma`. If False, `gamma` is not used.\n            When the next layer is linear (also e.g. `nn.relu`), this can be\n            disabled since the scaling will be done by the next layer.\n            Defaults to `True`.\n        rms_scaling: If True, `center` and `scale` are ignored, and the\n            inputs are scaled by `gamma` and the inverse square root\n            of the square of all inputs. This is an approximate and faster\n            approach that avoids ever computing the mean of the input.\n        beta_initializer: Initializer for the beta weight. Defaults to zeros.\n        gamma_initializer: Initializer for the gamma weight. Defaults to ones.\n        beta_regularizer: Optional regularizer for the beta weight.\n            None by default.\n        gamma_regularizer: Optional regularizer for the gamma weight.\n            None by default.\n        beta_constraint: Optional constraint for the beta weight.\n            None by default.\n        gamma_constraint: Optional constraint for the gamma weight.\n            None by default.\n        **kwargs: Base layer keyword arguments (e.g. `name` and `dtype`).\n\n\n    Reference:\n\n    - [Lei Ba et al., 2016](https://arxiv.org/abs/1607.06450).\n    \"\"\"\n\n    def __init__(\n        self,\n        axis=-1,\n        epsilon=1e-3,\n        center=True,\n        scale=True,\n        rms_scaling=False,\n        beta_initializer=\"zeros\",\n        gamma_initializer=\"ones\",\n        beta_regularizer=None,\n        gamma_regularizer=None,\n        beta_constraint=None,\n        gamma_constraint=None,\n        **kwargs\n    ):\n        super().__init__(**kwargs)\n        if isinstance(axis, (list, tuple)):\n            self.axis = list(axis)\n        elif isinstance(axis, int):\n            self.axis = axis\n        else:\n            raise TypeError(\n                \"Expected an int or a list/tuple of ints for the \"\n                \"argument 'axis', but received: %r\" % axis\n            )\n\n        self.epsilon = epsilon\n        self.center = center\n        self.scale = scale\n        self.rms_scaling = rms_scaling\n        self.beta_initializer = initializers.get(beta_initializer)\n        self.gamma_initializer = initializers.get(gamma_initializer)\n        self.beta_regularizer = regularizers.get(beta_regularizer)\n        self.gamma_regularizer = regularizers.get(gamma_regularizer)\n        self.beta_constraint = constraints.get(beta_constraint)\n        self.gamma_constraint = constraints.get(gamma_constraint)\n\n        self.supports_masking = True\n        self.autocast = False\n\n    def build(self, input_shape):\n        if isinstance(self.axis, list):\n            shape = tuple([input_shape[dim] for dim in self.axis])\n        else:\n            shape = (input_shape[self.axis],)\n            self.axis = [self.axis]\n        if self.scale or self.rms_scaling:\n            self.gamma = self.add_weight(\n                name=\"gamma\",\n                shape=shape,\n                initializer=self.gamma_initializer,\n                regularizer=self.gamma_regularizer,\n                constraint=self.gamma_constraint,\n                trainable=True,\n                autocast=False,\n            )\n        else:\n            self.gamma = None\n\n        if self.center and not self.rms_scaling:\n            self.beta = self.add_weight(\n                name=\"beta\",\n                shape=shape,\n                initializer=self.beta_initializer,\n                regularizer=self.beta_regularizer,\n                constraint=self.beta_constraint,\n                trainable=True,\n                autocast=False,\n            )\n        else:\n            self.beta = None\n\n        self.built = True\n\n    def call(self, inputs):\n        inputs = ops.cast(inputs, self.compute_dtype)\n        # Compute the axes along which to reduce the mean / variance\n        input_shape = inputs.shape\n        ndims = len(input_shape)\n\n        # Broadcasting only necessary for norm when the axis is not just\n        # the last dimension\n        broadcast_shape = [1] * ndims\n        for dim in self.axis:\n            broadcast_shape[dim] = input_shape[dim]\n\n        def _broadcast(v):\n            if (\n                v is not None\n                and len(v.shape) != ndims\n                and self.axis != [ndims - 1]\n            ):\n                return ops.reshape(v, broadcast_shape)\n            return v\n\n        input_dtype = inputs.dtype\n        if input_dtype in (\"float16\", \"bfloat16\") and self.dtype == \"float32\":\n            # If mixed precision is used, cast inputs to float32 so that\n            # this is at least as numerically stable as the fused version.\n            inputs = ops.cast(inputs, \"float32\")\n\n        if self.rms_scaling:\n            # Calculate outputs with only variance and gamma if rms scaling\n            # is enabled\n            # Calculate the variance along self.axis (layer activations).\n            variance = ops.var(inputs, axis=self.axis, keepdims=True)\n            inv = ops.rsqrt(variance + self.epsilon)\n\n            outputs = (\n                inputs * inv * ops.cast(_broadcast(self.gamma), inputs.dtype)\n            )\n        else:\n            # Calculate the mean & variance along self.axis (layer activations).\n            mean, variance = ops.moments(inputs, axes=self.axis, keepdims=True)\n            gamma, beta = _broadcast(self.gamma), _broadcast(self.beta)\n\n            inv = ops.rsqrt(variance + self.epsilon)\n            if gamma is not None:\n                gamma = ops.cast(gamma, inputs.dtype)\n                inv = inv * gamma\n\n            res = -mean * inv\n            if beta is not None:\n                beta = ops.cast(beta, inputs.dtype)\n                res = res + beta\n\n            outputs = inputs * inv + res\n\n        return ops.cast(outputs, input_dtype)\n\n    def compute_output_shape(self, input_shape):\n        return input_shape\n\n    def get_config(self):\n        config = {\n            \"axis\": self.axis,\n            \"epsilon\": self.epsilon,\n            \"center\": self.center,\n            \"scale\": self.scale,\n            \"beta_initializer\": initializers.serialize(self.beta_initializer),\n            \"gamma_initializer\": initializers.serialize(self.gamma_initializer),\n            \"beta_regularizer\": regularizers.serialize(self.beta_regularizer),\n            \"gamma_regularizer\": regularizers.serialize(self.gamma_regularizer),\n            \"beta_constraint\": constraints.serialize(self.beta_constraint),\n            \"gamma_constraint\": constraints.serialize(self.gamma_constraint),\n        }\n        base_config = super().get_config()\n        return {**base_config, **config}\n", "path": "keras/src/layers/normalization/layer_normalization.py"}]}