We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
tfa.layers.WeightNormalization fails with Assert when called by multiple threads
### Issue description

`tfa.layers.WeightNormalization` fails with the following Tensorflow assert failure when used by multiple threads (e.g., when the same layer wrapped by `tfa.layers.WeightNormalization` is part of two parallel execution flows in a Tensorflow graph while reusing the variables):
```
(0) Invalid argument: assertion failed: [The layer has been initialized.] [Condition x == y did not hold element-wise:] [x (discriminator/discriminator_1/n_layer_discriminator/sequential_21/grouped_conv1d_3/weight_normalization_304/cond/assert_equal_1/ReadVariableOp:0) = ] [1] [y (discriminator/discriminator_1/n_layer_discriminator/sequential_21/grouped_conv1d_3/weight_normalization_304/cond/assert_equal_1/y:0) = ] [0]
[[{{node discriminator/discriminator_1/n_layer_discriminator/sequential_21/grouped_conv1d_3/weight_normalization_304/cond/assert_equal_1/Assert/AssertGuard/Assert}}]]
```

### Proposed solution

Add the following line of code to the `tfa.layers.WeightNormalization.build` function:
```
self._initialized_critical_section = tf.CriticalSection(name='initialized_mutex')
```

Replace `g = tf.cond(self._initialized, _do_nothing, _update_weights)` in the `tfa.layers.WeightNormalization.call` function with the following lines of code:
```
def _get_initialized_weights():
    return tf.cond(self._initialized, _do_nothing, _update_weights)
g = self._initialized_critical_section.execute(_get_initialized_weights)
```

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `tensorflow_addons/layers/wrappers.py`
Content:
```
1 # Copyright 2019 The TensorFlow Authors. All Rights Reserved.
2 #
3 # Licensed under the Apache License, Version 2.0 (the "License");
4 # you may not use this file except in compliance with the License.
5 # You may obtain a copy of the License at
6 #
7 #     http://www.apache.org/licenses/LICENSE-2.0
8 #
9 # Unless required by applicable law or agreed to in writing, software
10 # distributed under the License is distributed on an "AS IS" BASIS,
11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 # See the License for the specific language governing permissions and
13 # limitations under the License.
14 # =============================================================================
15 from __future__ import absolute_import
16 from __future__ import division
17 from __future__ import print_function
18 
19 import tensorflow as tf
20 
21 
22 @tf.keras.utils.register_keras_serializable(package='Addons')
23 class WeightNormalization(tf.keras.layers.Wrapper):
24     """This wrapper reparameterizes a layer by decoupling the weight's
25     magnitude and direction.
26 
27     This speeds up convergence by improving the
28     conditioning of the optimization problem.
29     Weight Normalization: A Simple Reparameterization to Accelerate
30     Training of Deep Neural Networks: https://arxiv.org/abs/1602.07868
31     Tim Salimans, Diederik P. Kingma (2016)
32     WeightNormalization wrapper works for keras and tf layers.
33     ```python
34       net = WeightNormalization(
35           tf.keras.layers.Conv2D(2, 2, activation='relu'),
36           input_shape=(32, 32, 3),
37           data_init=True)(x)
38       net = WeightNormalization(
39           tf.keras.layers.Conv2D(16, 5, activation='relu'),
40           data_init=True)(net)
41       net = WeightNormalization(
42           tf.keras.layers.Dense(120, activation='relu'),
43           data_init=True)(net)
44       net = WeightNormalization(
45           tf.keras.layers.Dense(n_classes),
46           data_init=True)(net)
47     ```
48     Arguments:
49       layer: a layer instance.
50       data_init: If `True` use data dependent variable initialization
51     Raises:
52       ValueError: If not initialized with a `Layer` instance.
53       ValueError: If `Layer` does not contain a `kernel` of weights
54       NotImplementedError: If `data_init` is True and running graph execution
55     """
56 
57     def __init__(self, layer, data_init=True, **kwargs):
58         super(WeightNormalization, self).__init__(layer, **kwargs)
59         self.data_init = data_init
60         self._track_trackable(layer, name='layer')
61         self.is_rnn = isinstance(self.layer, tf.keras.layers.RNN)
62 
63     def build(self, input_shape):
64         """Build `Layer`"""
65         input_shape = tf.TensorShape(input_shape)
66         self.input_spec = tf.keras.layers.InputSpec(
67             shape=[None] + input_shape[1:])
68 
69         if not self.layer.built:
70             self.layer.build(input_shape)
71 
72         kernel_layer = self.layer.cell if self.is_rnn else self.layer
73 
74         if not hasattr(kernel_layer, 'kernel'):
75             raise ValueError('`WeightNormalization` must wrap a layer that'
76                              ' contains a `kernel` for weights')
77 
78         # The kernel's filter or unit dimension is -1
79         self.layer_depth = int(kernel_layer.kernel.shape[-1])
80         self.kernel_norm_axes = list(range(kernel_layer.kernel.shape.rank - 1))
81 
82         self.g = self.add_weight(
83             name='g',
84             shape=(self.layer_depth,),
85             initializer='ones',
86             dtype=kernel_layer.kernel.dtype,
87             trainable=True)
88         self.v = kernel_layer.kernel
89 
90         self._initialized = self.add_weight(
91             name='initialized',
92             shape=None,
93             initializer='zeros',
94             dtype=tf.dtypes.bool,
95             trainable=False)
96 
97         if self.data_init:
98             # Used for data initialization in self._data_dep_init.
99             with tf.name_scope('data_dep_init'):
100                 layer_config = tf.keras.layers.serialize(self.layer)
101                 layer_config['config']['trainable'] = False
102                 self._naked_clone_layer = tf.keras.layers.deserialize(
103                     layer_config)
104                 self._naked_clone_layer.build(input_shape)
105                 self._naked_clone_layer.set_weights(self.layer.get_weights())
106                 if self.is_rnn:
107                     self._naked_clone_layer.cell.activation = None
108                 else:
109                     self._naked_clone_layer.activation = None
110 
111         self.built = True
112 
113     def call(self, inputs):
114         """Call `Layer`"""
115 
116         def _do_nothing():
117             return tf.identity(self.g)
118 
119         def _update_weights():
120             # Ensure we read `self.g` after _update_weights.
121             with tf.control_dependencies(self._initialize_weights(inputs)):
122                 return tf.identity(self.g)
123 
124         g = tf.cond(self._initialized, _do_nothing, _update_weights)
125 
126         with tf.name_scope('compute_weights'):
127             # Replace kernel by normalized weight variable.
128             self.layer.kernel = tf.nn.l2_normalize(
129                 self.v, axis=self.kernel_norm_axes) * g
130 
131             # Ensure we calculate result after updating kernel.
132             update_kernel = tf.identity(self.layer.kernel)
133             with tf.control_dependencies([update_kernel]):
134                 outputs = self.layer(inputs)
135                 return outputs
136 
137     def compute_output_shape(self, input_shape):
138         return tf.TensorShape(
139             self.layer.compute_output_shape(input_shape).as_list())
140 
141     def _initialize_weights(self, inputs):
142         """Initialize weight g.
143 
144         The initial value of g could either from the initial value in v,
145         or by the input value if self.data_init is True.
146         """
147         with tf.control_dependencies([
148                 tf.debugging.assert_equal(  # pylint: disable=bad-continuation
149                     self._initialized,
150                     False,
151                     message='The layer has been initialized.')
152         ]):
153             if self.data_init:
154                 assign_tensors = self._data_dep_init(inputs)
155             else:
156                 assign_tensors = self._init_norm()
157             assign_tensors.append(self._initialized.assign(True))
158             return assign_tensors
159 
160     def _init_norm(self):
161         """Set the weight g with the norm of the weight vector."""
162         with tf.name_scope('init_norm'):
163             v_flat = tf.reshape(self.v, [-1, self.layer_depth])
164             v_norm = tf.linalg.norm(v_flat, axis=0)
165             g_tensor = self.g.assign(tf.reshape(v_norm, (self.layer_depth,)))
166             return [g_tensor]
167 
168     def _data_dep_init(self, inputs):
169         """Data dependent initialization."""
170         with tf.name_scope('data_dep_init'):
171             # Generate data dependent init values
172             x_init = self._naked_clone_layer(inputs)
173             data_norm_axes = list(range(x_init.shape.rank - 1))
174             m_init, v_init = tf.nn.moments(x_init, data_norm_axes)
175             scale_init = 1. / tf.math.sqrt(v_init + 1e-10)
176 
177             # Assign data dependent init values
178             g_tensor = self.g.assign(self.g * scale_init)
179             if hasattr(self.layer, 'bias') and self.layer.bias is not None:
180                 bias_tensor = self.layer.bias.assign(-m_init * scale_init)
181                 return [g_tensor, bias_tensor]
182             else:
183                 return [g_tensor]
184 
185     def get_config(self):
186         config = {'data_init': self.data_init}
187         base_config = super(WeightNormalization, self).get_config()
188         return dict(list(base_config.items()) + list(config.items()))
189
```
--- 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/tensorflow_addons/layers/wrappers.py b/tensorflow_addons/layers/wrappers.py
--- a/tensorflow_addons/layers/wrappers.py
+++ b/tensorflow_addons/layers/wrappers.py
@@ -58,6 +58,7 @@
         super(WeightNormalization, self).__init__(layer, **kwargs)
         self.data_init = data_init
         self._track_trackable(layer, name='layer')
+        self._init_critical_section = tf.CriticalSection(name='init_mutex')
         self.is_rnn = isinstance(self.layer, tf.keras.layers.RNN)
 
     def build(self, input_shape):
@@ -121,7 +122,8 @@
             with tf.control_dependencies(self._initialize_weights(inputs)):
                 return tf.identity(self.g)
 
-        g = tf.cond(self._initialized, _do_nothing, _update_weights)
+        g = self._init_critical_section.execute(lambda: tf.cond(
+            self._initialized, _do_nothing, _update_weights))
 
         with tf.name_scope('compute_weights'):
             # Replace kernel by normalized weight variable.

{"golden_diff": "diff --git a/tensorflow_addons/layers/wrappers.py b/tensorflow_addons/layers/wrappers.py\n--- a/tensorflow_addons/layers/wrappers.py\n+++ b/tensorflow_addons/layers/wrappers.py\n@@ -58,6 +58,7 @@\n         super(WeightNormalization, self).__init__(layer, **kwargs)\n         self.data_init = data_init\n         self._track_trackable(layer, name='layer')\n+        self._init_critical_section = tf.CriticalSection(name='init_mutex')\n         self.is_rnn = isinstance(self.layer, tf.keras.layers.RNN)\n \n     def build(self, input_shape):\n@@ -121,7 +122,8 @@\n             with tf.control_dependencies(self._initialize_weights(inputs)):\n                 return tf.identity(self.g)\n \n-        g = tf.cond(self._initialized, _do_nothing, _update_weights)\n+        g = self._init_critical_section.execute(lambda: tf.cond(\n+            self._initialized, _do_nothing, _update_weights))\n \n         with tf.name_scope('compute_weights'):\n             # Replace kernel by normalized weight variable.\n", "issue": "tfa.layers.WeightNormalization fails with Assert when called by multiple threads\n### Issue description\r\n\r\n`tfa.layers.WeightNormalization` fails with the following Tensorflow assert failure when used by multiple threads (e.g., when the same layer wrapped by `tfa.layers.WeightNormalization` is part of two parallel execution flows in a Tensorflow graph while reusing the variables):\r\n```\r\n(0) Invalid argument: assertion failed: [The layer has been initialized.] [Condition x == y did not hold element-wise:] [x (discriminator/discriminator_1/n_layer_discriminator/sequential_21/grouped_conv1d_3/weight_normalization_304/cond/assert_equal_1/ReadVariableOp:0) = ] [1] [y (discriminator/discriminator_1/n_layer_discriminator/sequential_21/grouped_conv1d_3/weight_normalization_304/cond/assert_equal_1/y:0) = ] [0]\r\n[[{{node discriminator/discriminator_1/n_layer_discriminator/sequential_21/grouped_conv1d_3/weight_normalization_304/cond/assert_equal_1/Assert/AssertGuard/Assert}}]]\r\n```\r\n\r\n### Proposed solution\r\n\r\nAdd the following line of code to the `tfa.layers.WeightNormalization.build` function:\r\n```\r\nself._initialized_critical_section = tf.CriticalSection(name='initialized_mutex')\r\n```\r\n\r\nReplace `g = tf.cond(self._initialized, _do_nothing, _update_weights)` in the `tfa.layers.WeightNormalization.call` function with the following lines of code:\r\n```\r\ndef _get_initialized_weights():\r\n    return tf.cond(self._initialized, _do_nothing, _update_weights)\r\ng = self._initialized_critical_section.execute(_get_initialized_weights)\r\n```\n", "before_files": [{"content": "# Copyright 2019 The TensorFlow Authors. All Rights Reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n# =============================================================================\nfrom __future__ import absolute_import\nfrom __future__ import division\nfrom __future__ import print_function\n\nimport tensorflow as tf\n\n\n@tf.keras.utils.register_keras_serializable(package='Addons')\nclass WeightNormalization(tf.keras.layers.Wrapper):\n    \"\"\"This wrapper reparameterizes a layer by decoupling the weight's\n    magnitude and direction.\n\n    This speeds up convergence by improving the\n    conditioning of the optimization problem.\n    Weight Normalization: A Simple Reparameterization to Accelerate\n    Training of Deep Neural Networks: https://arxiv.org/abs/1602.07868\n    Tim Salimans, Diederik P. Kingma (2016)\n    WeightNormalization wrapper works for keras and tf layers.\n    ```python\n      net = WeightNormalization(\n          tf.keras.layers.Conv2D(2, 2, activation='relu'),\n          input_shape=(32, 32, 3),\n          data_init=True)(x)\n      net = WeightNormalization(\n          tf.keras.layers.Conv2D(16, 5, activation='relu'),\n          data_init=True)(net)\n      net = WeightNormalization(\n          tf.keras.layers.Dense(120, activation='relu'),\n          data_init=True)(net)\n      net = WeightNormalization(\n          tf.keras.layers.Dense(n_classes),\n          data_init=True)(net)\n    ```\n    Arguments:\n      layer: a layer instance.\n      data_init: If `True` use data dependent variable initialization\n    Raises:\n      ValueError: If not initialized with a `Layer` instance.\n      ValueError: If `Layer` does not contain a `kernel` of weights\n      NotImplementedError: If `data_init` is True and running graph execution\n    \"\"\"\n\n    def __init__(self, layer, data_init=True, **kwargs):\n        super(WeightNormalization, self).__init__(layer, **kwargs)\n        self.data_init = data_init\n        self._track_trackable(layer, name='layer')\n        self.is_rnn = isinstance(self.layer, tf.keras.layers.RNN)\n\n    def build(self, input_shape):\n        \"\"\"Build `Layer`\"\"\"\n        input_shape = tf.TensorShape(input_shape)\n        self.input_spec = tf.keras.layers.InputSpec(\n            shape=[None] + input_shape[1:])\n\n        if not self.layer.built:\n            self.layer.build(input_shape)\n\n        kernel_layer = self.layer.cell if self.is_rnn else self.layer\n\n        if not hasattr(kernel_layer, 'kernel'):\n            raise ValueError('`WeightNormalization` must wrap a layer that'\n                             ' contains a `kernel` for weights')\n\n        # The kernel's filter or unit dimension is -1\n        self.layer_depth = int(kernel_layer.kernel.shape[-1])\n        self.kernel_norm_axes = list(range(kernel_layer.kernel.shape.rank - 1))\n\n        self.g = self.add_weight(\n            name='g',\n            shape=(self.layer_depth,),\n            initializer='ones',\n            dtype=kernel_layer.kernel.dtype,\n            trainable=True)\n        self.v = kernel_layer.kernel\n\n        self._initialized = self.add_weight(\n            name='initialized',\n            shape=None,\n            initializer='zeros',\n            dtype=tf.dtypes.bool,\n            trainable=False)\n\n        if self.data_init:\n            # Used for data initialization in self._data_dep_init.\n            with tf.name_scope('data_dep_init'):\n                layer_config = tf.keras.layers.serialize(self.layer)\n                layer_config['config']['trainable'] = False\n                self._naked_clone_layer = tf.keras.layers.deserialize(\n                    layer_config)\n                self._naked_clone_layer.build(input_shape)\n                self._naked_clone_layer.set_weights(self.layer.get_weights())\n                if self.is_rnn:\n                    self._naked_clone_layer.cell.activation = None\n                else:\n                    self._naked_clone_layer.activation = None\n\n        self.built = True\n\n    def call(self, inputs):\n        \"\"\"Call `Layer`\"\"\"\n\n        def _do_nothing():\n            return tf.identity(self.g)\n\n        def _update_weights():\n            # Ensure we read `self.g` after _update_weights.\n            with tf.control_dependencies(self._initialize_weights(inputs)):\n                return tf.identity(self.g)\n\n        g = tf.cond(self._initialized, _do_nothing, _update_weights)\n\n        with tf.name_scope('compute_weights'):\n            # Replace kernel by normalized weight variable.\n            self.layer.kernel = tf.nn.l2_normalize(\n                self.v, axis=self.kernel_norm_axes) * g\n\n            # Ensure we calculate result after updating kernel.\n            update_kernel = tf.identity(self.layer.kernel)\n            with tf.control_dependencies([update_kernel]):\n                outputs = self.layer(inputs)\n                return outputs\n\n    def compute_output_shape(self, input_shape):\n        return tf.TensorShape(\n            self.layer.compute_output_shape(input_shape).as_list())\n\n    def _initialize_weights(self, inputs):\n        \"\"\"Initialize weight g.\n\n        The initial value of g could either from the initial value in v,\n        or by the input value if self.data_init is True.\n        \"\"\"\n        with tf.control_dependencies([\n                tf.debugging.assert_equal(  # pylint: disable=bad-continuation\n                    self._initialized,\n                    False,\n                    message='The layer has been initialized.')\n        ]):\n            if self.data_init:\n                assign_tensors = self._data_dep_init(inputs)\n            else:\n                assign_tensors = self._init_norm()\n            assign_tensors.append(self._initialized.assign(True))\n            return assign_tensors\n\n    def _init_norm(self):\n        \"\"\"Set the weight g with the norm of the weight vector.\"\"\"\n        with tf.name_scope('init_norm'):\n            v_flat = tf.reshape(self.v, [-1, self.layer_depth])\n            v_norm = tf.linalg.norm(v_flat, axis=0)\n            g_tensor = self.g.assign(tf.reshape(v_norm, (self.layer_depth,)))\n            return [g_tensor]\n\n    def _data_dep_init(self, inputs):\n        \"\"\"Data dependent initialization.\"\"\"\n        with tf.name_scope('data_dep_init'):\n            # Generate data dependent init values\n            x_init = self._naked_clone_layer(inputs)\n            data_norm_axes = list(range(x_init.shape.rank - 1))\n            m_init, v_init = tf.nn.moments(x_init, data_norm_axes)\n            scale_init = 1. / tf.math.sqrt(v_init + 1e-10)\n\n            # Assign data dependent init values\n            g_tensor = self.g.assign(self.g * scale_init)\n            if hasattr(self.layer, 'bias') and self.layer.bias is not None:\n                bias_tensor = self.layer.bias.assign(-m_init * scale_init)\n                return [g_tensor, bias_tensor]\n            else:\n                return [g_tensor]\n\n    def get_config(self):\n        config = {'data_init': self.data_init}\n        base_config = super(WeightNormalization, self).get_config()\n        return dict(list(base_config.items()) + list(config.items()))\n", "path": "tensorflow_addons/layers/wrappers.py"}], "after_files": [{"content": "# Copyright 2019 The TensorFlow Authors. All Rights Reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n# =============================================================================\nfrom __future__ import absolute_import\nfrom __future__ import division\nfrom __future__ import print_function\n\nimport tensorflow as tf\n\n\n@tf.keras.utils.register_keras_serializable(package='Addons')\nclass WeightNormalization(tf.keras.layers.Wrapper):\n    \"\"\"This wrapper reparameterizes a layer by decoupling the weight's\n    magnitude and direction.\n\n    This speeds up convergence by improving the\n    conditioning of the optimization problem.\n    Weight Normalization: A Simple Reparameterization to Accelerate\n    Training of Deep Neural Networks: https://arxiv.org/abs/1602.07868\n    Tim Salimans, Diederik P. Kingma (2016)\n    WeightNormalization wrapper works for keras and tf layers.\n    ```python\n      net = WeightNormalization(\n          tf.keras.layers.Conv2D(2, 2, activation='relu'),\n          input_shape=(32, 32, 3),\n          data_init=True)(x)\n      net = WeightNormalization(\n          tf.keras.layers.Conv2D(16, 5, activation='relu'),\n          data_init=True)(net)\n      net = WeightNormalization(\n          tf.keras.layers.Dense(120, activation='relu'),\n          data_init=True)(net)\n      net = WeightNormalization(\n          tf.keras.layers.Dense(n_classes),\n          data_init=True)(net)\n    ```\n    Arguments:\n      layer: a layer instance.\n      data_init: If `True` use data dependent variable initialization\n    Raises:\n      ValueError: If not initialized with a `Layer` instance.\n      ValueError: If `Layer` does not contain a `kernel` of weights\n      NotImplementedError: If `data_init` is True and running graph execution\n    \"\"\"\n\n    def __init__(self, layer, data_init=True, **kwargs):\n        super(WeightNormalization, self).__init__(layer, **kwargs)\n        self.data_init = data_init\n        self._track_trackable(layer, name='layer')\n        self._init_critical_section = tf.CriticalSection(name='init_mutex')\n        self.is_rnn = isinstance(self.layer, tf.keras.layers.RNN)\n\n    def build(self, input_shape):\n        \"\"\"Build `Layer`\"\"\"\n        input_shape = tf.TensorShape(input_shape)\n        self.input_spec = tf.keras.layers.InputSpec(\n            shape=[None] + input_shape[1:])\n\n        if not self.layer.built:\n            self.layer.build(input_shape)\n\n        kernel_layer = self.layer.cell if self.is_rnn else self.layer\n\n        if not hasattr(kernel_layer, 'kernel'):\n            raise ValueError('`WeightNormalization` must wrap a layer that'\n                             ' contains a `kernel` for weights')\n\n        # The kernel's filter or unit dimension is -1\n        self.layer_depth = int(kernel_layer.kernel.shape[-1])\n        self.kernel_norm_axes = list(range(kernel_layer.kernel.shape.rank - 1))\n\n        self.g = self.add_weight(\n            name='g',\n            shape=(self.layer_depth,),\n            initializer='ones',\n            dtype=kernel_layer.kernel.dtype,\n            trainable=True)\n        self.v = kernel_layer.kernel\n\n        self._initialized = self.add_weight(\n            name='initialized',\n            shape=None,\n            initializer='zeros',\n            dtype=tf.dtypes.bool,\n            trainable=False)\n\n        if self.data_init:\n            # Used for data initialization in self._data_dep_init.\n            with tf.name_scope('data_dep_init'):\n                layer_config = tf.keras.layers.serialize(self.layer)\n                layer_config['config']['trainable'] = False\n                self._naked_clone_layer = tf.keras.layers.deserialize(\n                    layer_config)\n                self._naked_clone_layer.build(input_shape)\n                self._naked_clone_layer.set_weights(self.layer.get_weights())\n                if self.is_rnn:\n                    self._naked_clone_layer.cell.activation = None\n                else:\n                    self._naked_clone_layer.activation = None\n\n        self.built = True\n\n    def call(self, inputs):\n        \"\"\"Call `Layer`\"\"\"\n\n        def _do_nothing():\n            return tf.identity(self.g)\n\n        def _update_weights():\n            # Ensure we read `self.g` after _update_weights.\n            with tf.control_dependencies(self._initialize_weights(inputs)):\n                return tf.identity(self.g)\n\n        g = self._init_critical_section.execute(lambda: tf.cond(\n            self._initialized, _do_nothing, _update_weights))\n\n        with tf.name_scope('compute_weights'):\n            # Replace kernel by normalized weight variable.\n            self.layer.kernel = tf.nn.l2_normalize(\n                self.v, axis=self.kernel_norm_axes) * g\n\n            # Ensure we calculate result after updating kernel.\n            update_kernel = tf.identity(self.layer.kernel)\n            with tf.control_dependencies([update_kernel]):\n                outputs = self.layer(inputs)\n                return outputs\n\n    def compute_output_shape(self, input_shape):\n        return tf.TensorShape(\n            self.layer.compute_output_shape(input_shape).as_list())\n\n    def _initialize_weights(self, inputs):\n        \"\"\"Initialize weight g.\n\n        The initial value of g could either from the initial value in v,\n        or by the input value if self.data_init is True.\n        \"\"\"\n        with tf.control_dependencies([\n                tf.debugging.assert_equal(  # pylint: disable=bad-continuation\n                    self._initialized,\n                    False,\n                    message='The layer has been initialized.')\n        ]):\n            if self.data_init:\n                assign_tensors = self._data_dep_init(inputs)\n            else:\n                assign_tensors = self._init_norm()\n            assign_tensors.append(self._initialized.assign(True))\n            return assign_tensors\n\n    def _init_norm(self):\n        \"\"\"Set the weight g with the norm of the weight vector.\"\"\"\n        with tf.name_scope('init_norm'):\n            v_flat = tf.reshape(self.v, [-1, self.layer_depth])\n            v_norm = tf.linalg.norm(v_flat, axis=0)\n            g_tensor = self.g.assign(tf.reshape(v_norm, (self.layer_depth,)))\n            return [g_tensor]\n\n    def _data_dep_init(self, inputs):\n        \"\"\"Data dependent initialization.\"\"\"\n        with tf.name_scope('data_dep_init'):\n            # Generate data dependent init values\n            x_init = self._naked_clone_layer(inputs)\n            data_norm_axes = list(range(x_init.shape.rank - 1))\n            m_init, v_init = tf.nn.moments(x_init, data_norm_axes)\n            scale_init = 1. / tf.math.sqrt(v_init + 1e-10)\n\n            # Assign data dependent init values\n            g_tensor = self.g.assign(self.g * scale_init)\n            if hasattr(self.layer, 'bias') and self.layer.bias is not None:\n                bias_tensor = self.layer.bias.assign(-m_init * scale_init)\n                return [g_tensor, bias_tensor]\n            else:\n                return [g_tensor]\n\n    def get_config(self):\n        config = {'data_init': self.data_init}\n        base_config = super(WeightNormalization, self).get_config()\n        return dict(list(base_config.items()) + list(config.items()))\n", "path": "tensorflow_addons/layers/wrappers.py"}]}