We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
BatchNorm layer appears broken in multi-backend Keras
classic MNIST training:
- [Keras multi-backend with BN on JAX](https://colab.research.google.com/drive/1ZZI6UK9umgCDSIRcPqVUl6YpLmRdwMF3?resourcekey=0-8YS9usjjxUuh0FUl_JW0Ow&usp=sharing). Final accuracy __19%__!!!
- [Keras multi-backend w/o BN on JAX](https://colab.research.google.com/drive/16Yj63zZHmH2oQNBtNYmP5W0KJJ_RnLQi?resourcekey=0-Ut-aMYTyo3i36iJeQB4LRQ&usp=sharing). Final accuracy: 99%
- [Keras multi-backend with BN on TF](https://colab.research.google.com/drive/1ipSnQuGe2w4KGphG6hMeWdzt3H8m36cx?resourcekey=0-ZH-RRNE3rfXqqurX8yE_YA&usp=sharing). Final accuracy: __14%__!!!
- [Keras multi-backend w/o BN on TF](https://colab.research.google.com/drive/1beSnSJDvBwqYXzaGpvdVH9bUgCM2Jtsi?resourcekey=0-jmZOshPvFOUqWOEGq-z2Dg&usp=sharing). Final accuracy: 99%
- [Keras baseline (regular tf.keras on TF, with BN)](https://colab.research.google.com/drive/1dI3RTfVvYdYER09Q6onYIdsh3vkRmEci?resourcekey=0-FgN08JmB4pWDpsL8aePs8A&usp=sharing). Final accuracy: 99%

The common factors triggering the issue seem to be: using the BatchNorm layer, on Keras multi-backend. The choice of backend does not seem to change the behavior.

--- 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_core/layers/normalization/batch_normalization.py`
Content:
```
1 from keras_core import constraints
2 from keras_core import initializers
3 from keras_core import operations as ops
4 from keras_core import regularizers
5 from keras_core.api_export import keras_core_export
6 from keras_core.layers.input_spec import InputSpec
7 from keras_core.layers.layer import Layer
8 
9 
10 @keras_core_export("keras_core.layers.BatchNormalization")
11 class BatchNormalization(Layer):
12     """Layer that normalizes its inputs.
13 
14     Batch normalization applies a transformation that maintains the mean output
15     close to 0 and the output standard deviation close to 1.
16 
17     Importantly, batch normalization works differently during training and
18     during inference.
19 
20     **During training** (i.e. when using `fit()` or when calling the layer/model
21     with the argument `training=True`), the layer normalizes its output using
22     the mean and standard deviation of the current batch of inputs. That is to
23     say, for each channel being normalized, the layer returns
24     `gamma * (batch - mean(batch)) / sqrt(var(batch) + epsilon) + beta`, where:
25 
26     - `epsilon` is small constant (configurable as part of the constructor
27     arguments)
28     - `gamma` is a learned scaling factor (initialized as 1), which
29     can be disabled by passing `scale=False` to the constructor.
30     - `beta` is a learned offset factor (initialized as 0), which
31     can be disabled by passing `center=False` to the constructor.
32 
33     **During inference** (i.e. when using `evaluate()` or `predict()` or when
34     calling the layer/model with the argument `training=False` (which is the
35     default), the layer normalizes its output using a moving average of the
36     mean and standard deviation of the batches it has seen during training. That
37     is to say, it returns
38     `gamma * (batch - self.moving_mean) / sqrt(self.moving_var+epsilon) + beta`.
39 
40     `self.moving_mean` and `self.moving_var` are non-trainable variables that
41     are updated each time the layer in called in training mode, as such:
42 
43     - `moving_mean = moving_mean * momentum + mean(batch) * (1 - momentum)`
44     - `moving_var = moving_var * momentum + var(batch) * (1 - momentum)`
45 
46     As such, the layer will only normalize its inputs during inference
47     *after having been trained on data that has similar statistics as the
48     inference data*.
49 
50     Args:
51         axis: Integer, the axis that should be normalized
52             (typically the features axis). For instance, after a `Conv2D` layer
53             with `data_format="channels_first"`, use `axis=1`.
54         momentum: Momentum for the moving average.
55         epsilon: Small float added to variance to avoid dividing by zero.
56         center: If `True`, add offset of `beta` to normalized tensor.
57             If `False`, `beta` is ignored.
58         scale: If `True`, multiply by `gamma`. If `False`, `gamma` is not used.
59             When the next layer is linear this can be disabled
60             since the scaling will be done by the next layer.
61         beta_initializer: Initializer for the beta weight.
62         gamma_initializer: Initializer for the gamma weight.
63         moving_mean_initializer: Initializer for the moving mean.
64         moving_variance_initializer: Initializer for the moving variance.
65         beta_regularizer: Optional regularizer for the beta weight.
66         gamma_regularizer: Optional regularizer for the gamma weight.
67         beta_constraint: Optional constraint for the beta weight.
68         gamma_constraint: Optional constraint for the gamma weight.
69         **kwargs: Base layer keyword arguments (e.g. `name` and `dtype`).
70 
71     Call arguments:
72         inputs: Input tensor (of any rank).
73         training: Python boolean indicating whether the layer should behave in
74             training mode or in inference mode.
75             - `training=True`: The layer will normalize its inputs using
76             the mean and variance of the current batch of inputs.
77             - `training=False`: The layer will normalize its inputs using
78             the mean and variance of its moving statistics, learned during
79             training.
80 
81     Reference:
82 
83     - [Ioffe and Szegedy, 2015](https://arxiv.org/abs/1502.03167).
84 
85     **About setting `layer.trainable = False` on a `BatchNormalization` layer:**
86 
87     The meaning of setting `layer.trainable = False` is to freeze the layer,
88     i.e. its internal state will not change during training:
89     its trainable weights will not be updated
90     during `fit()` or `train_on_batch()`, and its state updates will not be run.
91 
92     Usually, this does not necessarily mean that the layer is run in inference
93     mode (which is normally controlled by the `training` argument that can
94     be passed when calling a layer). "Frozen state" and "inference mode"
95     are two separate concepts.
96 
97     However, in the case of the `BatchNormalization` layer, **setting
98     `trainable = False` on the layer means that the layer will be
99     subsequently run in inference mode** (meaning that it will use
100     the moving mean and the moving variance to normalize the current batch,
101     rather than using the mean and variance of the current batch).
102 
103     Note that:
104 
105     - Setting `trainable` on an model containing other layers will recursively
106         set the `trainable` value of all inner layers.
107     - If the value of the `trainable` attribute is changed after calling
108         `compile()` on a model, the new value doesn't take effect for this model
109         until `compile()` is called again.
110     """
111 
112     def __init__(
113         self,
114         axis=-1,
115         momentum=0.99,
116         epsilon=1e-3,
117         center=True,
118         scale=True,
119         beta_initializer="zeros",
120         gamma_initializer="ones",
121         moving_mean_initializer="zeros",
122         moving_variance_initializer="ones",
123         beta_regularizer=None,
124         gamma_regularizer=None,
125         beta_constraint=None,
126         gamma_constraint=None,
127         **kwargs,
128     ):
129         super().__init__(**kwargs)
130         self.axis = int(axis)
131         self.momentum = float(momentum)
132         self.epsilon = float(epsilon)
133         self.center = center
134         self.scale = scale
135         self.beta_initializer = initializers.get(beta_initializer)
136         self.gamma_initializer = initializers.get(gamma_initializer)
137         self.moving_mean_initializer = initializers.get(moving_mean_initializer)
138         self.moving_variance_initializer = initializers.get(
139             moving_variance_initializer
140         )
141         self.beta_regularizer = regularizers.get(beta_regularizer)
142         self.gamma_regularizer = regularizers.get(gamma_regularizer)
143         self.beta_constraint = constraints.get(beta_constraint)
144         self.gamma_constraint = constraints.get(gamma_constraint)
145         self.supports_masking = True
146 
147     def build(self, input_shape):
148         shape = (input_shape[self.axis],)
149         if self.scale:
150             self.gamma = self.add_weight(
151                 shape=shape,
152                 name="gamma",
153                 initializer=self.gamma_initializer,
154                 regularizer=self.gamma_regularizer,
155                 constraint=self.gamma_constraint,
156                 trainable=True,
157             )
158         if self.center:
159             self.beta = self.add_weight(
160                 shape=shape,
161                 name="beta",
162                 initializer=self.beta_initializer,
163                 regularizer=self.beta_regularizer,
164                 constraint=self.beta_constraint,
165                 trainable=True,
166             )
167         self.moving_mean = self.add_weight(
168             shape=shape,
169             name="moving_mean",
170             initializer=self.moving_mean_initializer,
171             trainable=False,
172         )
173         self.moving_variance = self.add_weight(
174             shape=shape,
175             name="moving_variance",
176             initializer=self.moving_variance_initializer,
177             trainable=False,
178         )
179         self.input_spec = InputSpec(
180             ndim=len(input_shape), axes={self.axis: input_shape[self.axis]}
181         )
182         reduction_axes = list(range(len(input_shape)))
183         del reduction_axes[self.axis]
184         self._reduction_axes = reduction_axes
185         self.built = True
186 
187     def compute_output_shape(self, input_shape):
188         return input_shape
189 
190     def call(self, inputs, training=None, mask=None):
191         broadcast_shape = [1] * len(inputs.shape)
192         broadcast_shape[self.axis] = inputs.shape[self.axis]
193         if training and self.trainable:
194             mean = ops.mean(inputs, axis=self._reduction_axes, keepdims=True)
195             variance = ops.mean(
196                 ops.square(inputs), axis=self._reduction_axes, keepdims=True
197             ) - ops.square(mean)
198             outputs = (inputs - mean) / ops.sqrt(variance + self.epsilon)
199             mean = ops.squeeze(mean, self._reduction_axes)
200             variance = ops.squeeze(variance, self._reduction_axes)
201             self.moving_mean.assign(
202                 self.moving_mean * self.momentum + mean * (1.0 - self.momentum)
203             )
204         else:
205             moving_mean = ops.reshape(self.moving_mean, broadcast_shape)
206             moving_variance = ops.reshape(self.moving_variance, broadcast_shape)
207             outputs = (inputs - moving_mean) / ops.sqrt(
208                 moving_variance + self.epsilon
209             )
210         if self.scale:
211             gamma = ops.reshape(self.gamma, broadcast_shape)
212             outputs = outputs * gamma
213         if self.center:
214             beta = ops.reshape(self.beta, broadcast_shape)
215             outputs = outputs + beta
216         return outputs
217 
218     def get_config(self):
219         base_config = super().get_config()
220         config = {
221             "axis": self.axis,
222             "momentum": self.momentum,
223             "epsilon": self.epsilon,
224             "center": self.center,
225             "scale": self.scale,
226             "beta_initializer": initializers.serialize(self.beta_initializer),
227             "gamma_initializer": initializers.serialize(self.gamma_initializer),
228             "moving_mean_initializer": initializers.serialize(
229                 self.moving_mean_initializer
230             ),
231             "moving_variance_initializer": initializers.serialize(
232                 self.moving_variance_initializer
233             ),
234             "beta_regularizer": regularizers.serialize(self.beta_regularizer),
235             "gamma_regularizer": regularizers.serialize(self.gamma_regularizer),
236             "beta_constraint": constraints.serialize(self.beta_constraint),
237             "gamma_constraint": constraints.serialize(self.gamma_constraint),
238         }
239         return {**base_config, **config}
240
```
--- 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_core/layers/normalization/batch_normalization.py b/keras_core/layers/normalization/batch_normalization.py
--- a/keras_core/layers/normalization/batch_normalization.py
+++ b/keras_core/layers/normalization/batch_normalization.py
@@ -201,6 +201,10 @@
             self.moving_mean.assign(
                 self.moving_mean * self.momentum + mean * (1.0 - self.momentum)
             )
+            self.moving_variance.assign(
+                self.moving_variance * self.momentum
+                + variance * (1.0 - self.momentum)
+            )
         else:
             moving_mean = ops.reshape(self.moving_mean, broadcast_shape)
             moving_variance = ops.reshape(self.moving_variance, broadcast_shape)

{"golden_diff": "diff --git a/keras_core/layers/normalization/batch_normalization.py b/keras_core/layers/normalization/batch_normalization.py\n--- a/keras_core/layers/normalization/batch_normalization.py\n+++ b/keras_core/layers/normalization/batch_normalization.py\n@@ -201,6 +201,10 @@\n             self.moving_mean.assign(\n                 self.moving_mean * self.momentum + mean * (1.0 - self.momentum)\n             )\n+            self.moving_variance.assign(\n+                self.moving_variance * self.momentum\n+                + variance * (1.0 - self.momentum)\n+            )\n         else:\n             moving_mean = ops.reshape(self.moving_mean, broadcast_shape)\n             moving_variance = ops.reshape(self.moving_variance, broadcast_shape)\n", "issue": "BatchNorm layer appears broken in multi-backend Keras\nclassic MNIST training:\r\n- [Keras multi-backend with BN on JAX](https://colab.research.google.com/drive/1ZZI6UK9umgCDSIRcPqVUl6YpLmRdwMF3?resourcekey=0-8YS9usjjxUuh0FUl_JW0Ow&usp=sharing). Final accuracy __19%__!!!\r\n- [Keras multi-backend w/o BN on JAX](https://colab.research.google.com/drive/16Yj63zZHmH2oQNBtNYmP5W0KJJ_RnLQi?resourcekey=0-Ut-aMYTyo3i36iJeQB4LRQ&usp=sharing). Final accuracy: 99%\r\n- [Keras multi-backend with BN on TF](https://colab.research.google.com/drive/1ipSnQuGe2w4KGphG6hMeWdzt3H8m36cx?resourcekey=0-ZH-RRNE3rfXqqurX8yE_YA&usp=sharing). Final accuracy: __14%__!!!\r\n- [Keras multi-backend w/o BN on TF](https://colab.research.google.com/drive/1beSnSJDvBwqYXzaGpvdVH9bUgCM2Jtsi?resourcekey=0-jmZOshPvFOUqWOEGq-z2Dg&usp=sharing). Final accuracy: 99%\r\n- [Keras baseline (regular tf.keras on TF, with BN)](https://colab.research.google.com/drive/1dI3RTfVvYdYER09Q6onYIdsh3vkRmEci?resourcekey=0-FgN08JmB4pWDpsL8aePs8A&usp=sharing). Final accuracy: 99%\r\n\r\nThe common factors triggering the issue seem to be: using the BatchNorm layer, on Keras multi-backend. The choice of backend does not seem to change the behavior.\n", "before_files": [{"content": "from keras_core import constraints\nfrom keras_core import initializers\nfrom keras_core import operations as ops\nfrom keras_core import regularizers\nfrom keras_core.api_export import keras_core_export\nfrom keras_core.layers.input_spec import InputSpec\nfrom keras_core.layers.layer import Layer\n\n\n@keras_core_export(\"keras_core.layers.BatchNormalization\")\nclass BatchNormalization(Layer):\n    \"\"\"Layer that normalizes its inputs.\n\n    Batch normalization applies a transformation that maintains the mean output\n    close to 0 and the output standard deviation close to 1.\n\n    Importantly, batch normalization works differently during training and\n    during inference.\n\n    **During training** (i.e. when using `fit()` or when calling the layer/model\n    with the argument `training=True`), the layer normalizes its output using\n    the mean and standard deviation of the current batch of inputs. That is to\n    say, for each channel being normalized, the layer returns\n    `gamma * (batch - mean(batch)) / sqrt(var(batch) + epsilon) + beta`, where:\n\n    - `epsilon` is small constant (configurable as part of the constructor\n    arguments)\n    - `gamma` is a learned scaling factor (initialized as 1), which\n    can be disabled by passing `scale=False` to the constructor.\n    - `beta` is a learned offset factor (initialized as 0), which\n    can be disabled by passing `center=False` to the constructor.\n\n    **During inference** (i.e. when using `evaluate()` or `predict()` or when\n    calling the layer/model with the argument `training=False` (which is the\n    default), the layer normalizes its output using a moving average of the\n    mean and standard deviation of the batches it has seen during training. That\n    is to say, it returns\n    `gamma * (batch - self.moving_mean) / sqrt(self.moving_var+epsilon) + beta`.\n\n    `self.moving_mean` and `self.moving_var` are non-trainable variables that\n    are updated each time the layer in called in training mode, as such:\n\n    - `moving_mean = moving_mean * momentum + mean(batch) * (1 - momentum)`\n    - `moving_var = moving_var * momentum + var(batch) * (1 - momentum)`\n\n    As such, the layer will only normalize its inputs during inference\n    *after having been trained on data that has similar statistics as the\n    inference data*.\n\n    Args:\n        axis: Integer, the axis that should be normalized\n            (typically the features axis). For instance, after a `Conv2D` layer\n            with `data_format=\"channels_first\"`, use `axis=1`.\n        momentum: Momentum for the moving average.\n        epsilon: Small float added to variance to avoid dividing by zero.\n        center: If `True`, add offset of `beta` to normalized tensor.\n            If `False`, `beta` is ignored.\n        scale: If `True`, multiply by `gamma`. If `False`, `gamma` is not used.\n            When the next layer is linear this can be disabled\n            since the scaling will be done by the next layer.\n        beta_initializer: Initializer for the beta weight.\n        gamma_initializer: Initializer for the gamma weight.\n        moving_mean_initializer: Initializer for the moving mean.\n        moving_variance_initializer: Initializer for the moving variance.\n        beta_regularizer: Optional regularizer for the beta weight.\n        gamma_regularizer: Optional regularizer for the gamma weight.\n        beta_constraint: Optional constraint for the beta weight.\n        gamma_constraint: Optional constraint for the gamma weight.\n        **kwargs: Base layer keyword arguments (e.g. `name` and `dtype`).\n\n    Call arguments:\n        inputs: Input tensor (of any rank).\n        training: Python boolean indicating whether the layer should behave in\n            training mode or in inference mode.\n            - `training=True`: The layer will normalize its inputs using\n            the mean and variance of the current batch of inputs.\n            - `training=False`: The layer will normalize its inputs using\n            the mean and variance of its moving statistics, learned during\n            training.\n\n    Reference:\n\n    - [Ioffe and Szegedy, 2015](https://arxiv.org/abs/1502.03167).\n\n    **About setting `layer.trainable = False` on a `BatchNormalization` layer:**\n\n    The meaning of setting `layer.trainable = False` is to freeze the layer,\n    i.e. its internal state will not change during training:\n    its trainable weights will not be updated\n    during `fit()` or `train_on_batch()`, and its state updates will not be run.\n\n    Usually, this does not necessarily mean that the layer is run in inference\n    mode (which is normally controlled by the `training` argument that can\n    be passed when calling a layer). \"Frozen state\" and \"inference mode\"\n    are two separate concepts.\n\n    However, in the case of the `BatchNormalization` layer, **setting\n    `trainable = False` on the layer means that the layer will be\n    subsequently run in inference mode** (meaning that it will use\n    the moving mean and the moving variance to normalize the current batch,\n    rather than using the mean and variance of the current batch).\n\n    Note that:\n\n    - Setting `trainable` on an model containing other layers will recursively\n        set the `trainable` value of all inner layers.\n    - If the value of the `trainable` attribute is changed after calling\n        `compile()` on a model, the new value doesn't take effect for this model\n        until `compile()` is called again.\n    \"\"\"\n\n    def __init__(\n        self,\n        axis=-1,\n        momentum=0.99,\n        epsilon=1e-3,\n        center=True,\n        scale=True,\n        beta_initializer=\"zeros\",\n        gamma_initializer=\"ones\",\n        moving_mean_initializer=\"zeros\",\n        moving_variance_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        self.axis = int(axis)\n        self.momentum = float(momentum)\n        self.epsilon = float(epsilon)\n        self.center = center\n        self.scale = scale\n        self.beta_initializer = initializers.get(beta_initializer)\n        self.gamma_initializer = initializers.get(gamma_initializer)\n        self.moving_mean_initializer = initializers.get(moving_mean_initializer)\n        self.moving_variance_initializer = initializers.get(\n            moving_variance_initializer\n        )\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        self.supports_masking = True\n\n    def build(self, input_shape):\n        shape = (input_shape[self.axis],)\n        if self.scale:\n            self.gamma = self.add_weight(\n                shape=shape,\n                name=\"gamma\",\n                initializer=self.gamma_initializer,\n                regularizer=self.gamma_regularizer,\n                constraint=self.gamma_constraint,\n                trainable=True,\n            )\n        if self.center:\n            self.beta = self.add_weight(\n                shape=shape,\n                name=\"beta\",\n                initializer=self.beta_initializer,\n                regularizer=self.beta_regularizer,\n                constraint=self.beta_constraint,\n                trainable=True,\n            )\n        self.moving_mean = self.add_weight(\n            shape=shape,\n            name=\"moving_mean\",\n            initializer=self.moving_mean_initializer,\n            trainable=False,\n        )\n        self.moving_variance = self.add_weight(\n            shape=shape,\n            name=\"moving_variance\",\n            initializer=self.moving_variance_initializer,\n            trainable=False,\n        )\n        self.input_spec = InputSpec(\n            ndim=len(input_shape), axes={self.axis: input_shape[self.axis]}\n        )\n        reduction_axes = list(range(len(input_shape)))\n        del reduction_axes[self.axis]\n        self._reduction_axes = reduction_axes\n        self.built = True\n\n    def compute_output_shape(self, input_shape):\n        return input_shape\n\n    def call(self, inputs, training=None, mask=None):\n        broadcast_shape = [1] * len(inputs.shape)\n        broadcast_shape[self.axis] = inputs.shape[self.axis]\n        if training and self.trainable:\n            mean = ops.mean(inputs, axis=self._reduction_axes, keepdims=True)\n            variance = ops.mean(\n                ops.square(inputs), axis=self._reduction_axes, keepdims=True\n            ) - ops.square(mean)\n            outputs = (inputs - mean) / ops.sqrt(variance + self.epsilon)\n            mean = ops.squeeze(mean, self._reduction_axes)\n            variance = ops.squeeze(variance, self._reduction_axes)\n            self.moving_mean.assign(\n                self.moving_mean * self.momentum + mean * (1.0 - self.momentum)\n            )\n        else:\n            moving_mean = ops.reshape(self.moving_mean, broadcast_shape)\n            moving_variance = ops.reshape(self.moving_variance, broadcast_shape)\n            outputs = (inputs - moving_mean) / ops.sqrt(\n                moving_variance + self.epsilon\n            )\n        if self.scale:\n            gamma = ops.reshape(self.gamma, broadcast_shape)\n            outputs = outputs * gamma\n        if self.center:\n            beta = ops.reshape(self.beta, broadcast_shape)\n            outputs = outputs + beta\n        return outputs\n\n    def get_config(self):\n        base_config = super().get_config()\n        config = {\n            \"axis\": self.axis,\n            \"momentum\": self.momentum,\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            \"moving_mean_initializer\": initializers.serialize(\n                self.moving_mean_initializer\n            ),\n            \"moving_variance_initializer\": initializers.serialize(\n                self.moving_variance_initializer\n            ),\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        return {**base_config, **config}\n", "path": "keras_core/layers/normalization/batch_normalization.py"}], "after_files": [{"content": "from keras_core import constraints\nfrom keras_core import initializers\nfrom keras_core import operations as ops\nfrom keras_core import regularizers\nfrom keras_core.api_export import keras_core_export\nfrom keras_core.layers.input_spec import InputSpec\nfrom keras_core.layers.layer import Layer\n\n\n@keras_core_export(\"keras_core.layers.BatchNormalization\")\nclass BatchNormalization(Layer):\n    \"\"\"Layer that normalizes its inputs.\n\n    Batch normalization applies a transformation that maintains the mean output\n    close to 0 and the output standard deviation close to 1.\n\n    Importantly, batch normalization works differently during training and\n    during inference.\n\n    **During training** (i.e. when using `fit()` or when calling the layer/model\n    with the argument `training=True`), the layer normalizes its output using\n    the mean and standard deviation of the current batch of inputs. That is to\n    say, for each channel being normalized, the layer returns\n    `gamma * (batch - mean(batch)) / sqrt(var(batch) + epsilon) + beta`, where:\n\n    - `epsilon` is small constant (configurable as part of the constructor\n    arguments)\n    - `gamma` is a learned scaling factor (initialized as 1), which\n    can be disabled by passing `scale=False` to the constructor.\n    - `beta` is a learned offset factor (initialized as 0), which\n    can be disabled by passing `center=False` to the constructor.\n\n    **During inference** (i.e. when using `evaluate()` or `predict()` or when\n    calling the layer/model with the argument `training=False` (which is the\n    default), the layer normalizes its output using a moving average of the\n    mean and standard deviation of the batches it has seen during training. That\n    is to say, it returns\n    `gamma * (batch - self.moving_mean) / sqrt(self.moving_var+epsilon) + beta`.\n\n    `self.moving_mean` and `self.moving_var` are non-trainable variables that\n    are updated each time the layer in called in training mode, as such:\n\n    - `moving_mean = moving_mean * momentum + mean(batch) * (1 - momentum)`\n    - `moving_var = moving_var * momentum + var(batch) * (1 - momentum)`\n\n    As such, the layer will only normalize its inputs during inference\n    *after having been trained on data that has similar statistics as the\n    inference data*.\n\n    Args:\n        axis: Integer, the axis that should be normalized\n            (typically the features axis). For instance, after a `Conv2D` layer\n            with `data_format=\"channels_first\"`, use `axis=1`.\n        momentum: Momentum for the moving average.\n        epsilon: Small float added to variance to avoid dividing by zero.\n        center: If `True`, add offset of `beta` to normalized tensor.\n            If `False`, `beta` is ignored.\n        scale: If `True`, multiply by `gamma`. If `False`, `gamma` is not used.\n            When the next layer is linear this can be disabled\n            since the scaling will be done by the next layer.\n        beta_initializer: Initializer for the beta weight.\n        gamma_initializer: Initializer for the gamma weight.\n        moving_mean_initializer: Initializer for the moving mean.\n        moving_variance_initializer: Initializer for the moving variance.\n        beta_regularizer: Optional regularizer for the beta weight.\n        gamma_regularizer: Optional regularizer for the gamma weight.\n        beta_constraint: Optional constraint for the beta weight.\n        gamma_constraint: Optional constraint for the gamma weight.\n        **kwargs: Base layer keyword arguments (e.g. `name` and `dtype`).\n\n    Call arguments:\n        inputs: Input tensor (of any rank).\n        training: Python boolean indicating whether the layer should behave in\n            training mode or in inference mode.\n            - `training=True`: The layer will normalize its inputs using\n            the mean and variance of the current batch of inputs.\n            - `training=False`: The layer will normalize its inputs using\n            the mean and variance of its moving statistics, learned during\n            training.\n\n    Reference:\n\n    - [Ioffe and Szegedy, 2015](https://arxiv.org/abs/1502.03167).\n\n    **About setting `layer.trainable = False` on a `BatchNormalization` layer:**\n\n    The meaning of setting `layer.trainable = False` is to freeze the layer,\n    i.e. its internal state will not change during training:\n    its trainable weights will not be updated\n    during `fit()` or `train_on_batch()`, and its state updates will not be run.\n\n    Usually, this does not necessarily mean that the layer is run in inference\n    mode (which is normally controlled by the `training` argument that can\n    be passed when calling a layer). \"Frozen state\" and \"inference mode\"\n    are two separate concepts.\n\n    However, in the case of the `BatchNormalization` layer, **setting\n    `trainable = False` on the layer means that the layer will be\n    subsequently run in inference mode** (meaning that it will use\n    the moving mean and the moving variance to normalize the current batch,\n    rather than using the mean and variance of the current batch).\n\n    Note that:\n\n    - Setting `trainable` on an model containing other layers will recursively\n        set the `trainable` value of all inner layers.\n    - If the value of the `trainable` attribute is changed after calling\n        `compile()` on a model, the new value doesn't take effect for this model\n        until `compile()` is called again.\n    \"\"\"\n\n    def __init__(\n        self,\n        axis=-1,\n        momentum=0.99,\n        epsilon=1e-3,\n        center=True,\n        scale=True,\n        beta_initializer=\"zeros\",\n        gamma_initializer=\"ones\",\n        moving_mean_initializer=\"zeros\",\n        moving_variance_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        self.axis = int(axis)\n        self.momentum = float(momentum)\n        self.epsilon = float(epsilon)\n        self.center = center\n        self.scale = scale\n        self.beta_initializer = initializers.get(beta_initializer)\n        self.gamma_initializer = initializers.get(gamma_initializer)\n        self.moving_mean_initializer = initializers.get(moving_mean_initializer)\n        self.moving_variance_initializer = initializers.get(\n            moving_variance_initializer\n        )\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        self.supports_masking = True\n\n    def build(self, input_shape):\n        shape = (input_shape[self.axis],)\n        if self.scale:\n            self.gamma = self.add_weight(\n                shape=shape,\n                name=\"gamma\",\n                initializer=self.gamma_initializer,\n                regularizer=self.gamma_regularizer,\n                constraint=self.gamma_constraint,\n                trainable=True,\n            )\n        if self.center:\n            self.beta = self.add_weight(\n                shape=shape,\n                name=\"beta\",\n                initializer=self.beta_initializer,\n                regularizer=self.beta_regularizer,\n                constraint=self.beta_constraint,\n                trainable=True,\n            )\n        self.moving_mean = self.add_weight(\n            shape=shape,\n            name=\"moving_mean\",\n            initializer=self.moving_mean_initializer,\n            trainable=False,\n        )\n        self.moving_variance = self.add_weight(\n            shape=shape,\n            name=\"moving_variance\",\n            initializer=self.moving_variance_initializer,\n            trainable=False,\n        )\n        self.input_spec = InputSpec(\n            ndim=len(input_shape), axes={self.axis: input_shape[self.axis]}\n        )\n        reduction_axes = list(range(len(input_shape)))\n        del reduction_axes[self.axis]\n        self._reduction_axes = reduction_axes\n        self.built = True\n\n    def compute_output_shape(self, input_shape):\n        return input_shape\n\n    def call(self, inputs, training=None, mask=None):\n        broadcast_shape = [1] * len(inputs.shape)\n        broadcast_shape[self.axis] = inputs.shape[self.axis]\n        if training and self.trainable:\n            mean = ops.mean(inputs, axis=self._reduction_axes, keepdims=True)\n            variance = ops.mean(\n                ops.square(inputs), axis=self._reduction_axes, keepdims=True\n            ) - ops.square(mean)\n            outputs = (inputs - mean) / ops.sqrt(variance + self.epsilon)\n            mean = ops.squeeze(mean, self._reduction_axes)\n            variance = ops.squeeze(variance, self._reduction_axes)\n            self.moving_mean.assign(\n                self.moving_mean * self.momentum + mean * (1.0 - self.momentum)\n            )\n            self.moving_variance.assign(\n                self.moving_variance * self.momentum\n                + variance * (1.0 - self.momentum)\n            )\n        else:\n            moving_mean = ops.reshape(self.moving_mean, broadcast_shape)\n            moving_variance = ops.reshape(self.moving_variance, broadcast_shape)\n            outputs = (inputs - moving_mean) / ops.sqrt(\n                moving_variance + self.epsilon\n            )\n        if self.scale:\n            gamma = ops.reshape(self.gamma, broadcast_shape)\n            outputs = outputs * gamma\n        if self.center:\n            beta = ops.reshape(self.beta, broadcast_shape)\n            outputs = outputs + beta\n        return outputs\n\n    def get_config(self):\n        base_config = super().get_config()\n        config = {\n            \"axis\": self.axis,\n            \"momentum\": self.momentum,\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            \"moving_mean_initializer\": initializers.serialize(\n                self.moving_mean_initializer\n            ),\n            \"moving_variance_initializer\": initializers.serialize(\n                self.moving_variance_initializer\n            ),\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        return {**base_config, **config}\n", "path": "keras_core/layers/normalization/batch_normalization.py"}]}