We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
WeightNormalization fails after update
Yesterday I cut a 0.3 release on pypi test, but as a sanity check on the weightnorm changes #208 I ran a quick regression check:

https://colab.research.google.com/drive/1RQlVR9X7vj8q3W3sJC6YA3Cq45XTXrau

The new layer is failing with loss of nan. Looking at the test coverage for this layer I'm struggling how to best catch this without running a full model. The losses for our toy test cases do not return nan so perhaps we need a true regression test like mnist?

EDIT -- Fix colab link

--- 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 from tensorflow_addons.utils import keras_utils
21 
22 
23 @keras_utils.register_keras_custom_object
24 class WeightNormalization(tf.keras.layers.Wrapper):
25     """This wrapper reparameterizes a layer by decoupling the weight's
26     magnitude and direction.
27 
28     This speeds up convergence by improving the
29     conditioning of the optimization problem.
30     Weight Normalization: A Simple Reparameterization to Accelerate
31     Training of Deep Neural Networks: https://arxiv.org/abs/1602.07868
32     Tim Salimans, Diederik P. Kingma (2016)
33     WeightNormalization wrapper works for keras and tf layers.
34     ```python
35       net = WeightNormalization(
36           tf.keras.layers.Conv2D(2, 2, activation='relu'),
37           input_shape=(32, 32, 3),
38           data_init=True)(x)
39       net = WeightNormalization(
40           tf.keras.layers.Conv2D(16, 5, activation='relu'),
41           data_init=True)(net)
42       net = WeightNormalization(
43           tf.keras.layers.Dense(120, activation='relu'),
44           data_init=True)(net)
45       net = WeightNormalization(
46           tf.keras.layers.Dense(n_classes),
47           data_init=True)(net)
48     ```
49     Arguments:
50       layer: a layer instance.
51       data_init: If `True` use data dependent variable initialization
52     Raises:
53       ValueError: If not initialized with a `Layer` instance.
54       ValueError: If `Layer` does not contain a `kernel` of weights
55       NotImplementedError: If `data_init` is True and running graph execution
56     """
57 
58     def __init__(self, layer, data_init=True, **kwargs):
59         super(WeightNormalization, self).__init__(layer, **kwargs)
60         self.data_init = data_init
61         self._initialized = False
62         self._track_trackable(layer, name='layer')
63 
64     def build(self, input_shape):
65         """Build `Layer`"""
66         input_shape = tf.TensorShape(input_shape).as_list()
67         self.input_spec = tf.keras.layers.InputSpec(shape=input_shape)
68 
69         if not self.layer.built:
70             self.layer.build(input_shape)
71 
72             if not hasattr(self.layer, 'kernel'):
73                 raise ValueError('`WeightNormalization` must wrap a layer that'
74                                  ' contains a `kernel` for weights')
75 
76             # The kernel's filter or unit dimension is -1
77             self.layer_depth = int(self.layer.kernel.shape[-1])
78             self.kernel_norm_axes = list(
79                 range(self.layer.kernel.shape.rank - 1))
80 
81             self.v = self.layer.kernel
82             self.g = self.add_variable(
83                 name="g",
84                 shape=(self.layer_depth,),
85                 initializer=tf.keras.initializers.get('ones'),
86                 dtype=self.layer.kernel.dtype,
87                 trainable=True)
88 
89         super(WeightNormalization, self).build()
90 
91     def call(self, inputs):
92         """Call `Layer`"""
93         if not self._initialized:
94             self._initialize_weights(inputs)
95 
96         self._compute_weights()  # Recompute weights for each forward pass
97         output = self.layer(inputs)
98         return output
99 
100     def compute_output_shape(self, input_shape):
101         return tf.TensorShape(
102             self.layer.compute_output_shape(input_shape).as_list())
103 
104     def _compute_weights(self):
105         """Generate normalized weights.
106 
107         This method will update the value of self.layer.kernel with the
108         normalized value, so that the layer is ready for call().
109         """
110         with tf.name_scope('compute_weights'):
111             self.layer.kernel = tf.nn.l2_normalize(
112                 self.v, axis=self.kernel_norm_axes) * self.g
113 
114     def _initialize_weights(self, inputs):
115         """Initialize weight g.
116 
117         The initial value of g could either from the initial value in v,
118         or by the input value if self.data_init is True.
119         """
120         if self.data_init:
121             self._data_dep_init(inputs)
122         else:
123             self._init_norm()
124         self._initialized = True
125 
126     def _init_norm(self):
127         """Set the weight g with the norm of the weight vector."""
128         with tf.name_scope('init_norm'):
129             flat = tf.reshape(self.v, [-1, self.layer_depth])
130             self.g.assign(
131                 tf.reshape(tf.linalg.norm(flat, axis=0), (self.layer_depth,)))
132 
133     def _data_dep_init(self, inputs):
134         """Data dependent initialization."""
135 
136         with tf.name_scope('data_dep_init'):
137             # Generate data dependent init values
138             existing_activation = self.layer.activation
139             self.layer.activation = None
140             x_init = self.layer(inputs)
141             data_norm_axes = list(range(x_init.shape.rank - 1))
142             m_init, v_init = tf.nn.moments(x_init, data_norm_axes)
143             scale_init = 1. / tf.math.sqrt(v_init + 1e-10)
144 
145         # Assign data dependent init values
146         self.g.assign(self.g * scale_init)
147         if hasattr(self.layer, 'bias'):
148             self.layer.bias.assign(-m_init * scale_init)
149         self.layer.activation = existing_activation
150 
151     def get_config(self):
152         config = {'data_init': self.data_init}
153         base_config = super(WeightNormalization, self).get_config()
154         return dict(list(base_config.items()) + list(config.items()))
155
```
--- 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
@@ -88,6 +88,7 @@
 
         super(WeightNormalization, self).build()
 
+    @tf.function
     def call(self, inputs):
         """Call `Layer`"""
         if not self._initialized:
@@ -143,9 +144,9 @@
             scale_init = 1. / tf.math.sqrt(v_init + 1e-10)
 
         # Assign data dependent init values
-        self.g.assign(self.g * scale_init)
+        self.g = self.g * scale_init
         if hasattr(self.layer, 'bias'):
-            self.layer.bias.assign(-m_init * scale_init)
+            self.layer.bias = -m_init * scale_init
         self.layer.activation = existing_activation
 
     def get_config(self):

{"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@@ -88,6 +88,7 @@\n \n         super(WeightNormalization, self).build()\n \n+    @tf.function\n     def call(self, inputs):\n         \"\"\"Call `Layer`\"\"\"\n         if not self._initialized:\n@@ -143,9 +144,9 @@\n             scale_init = 1. / tf.math.sqrt(v_init + 1e-10)\n \n         # Assign data dependent init values\n-        self.g.assign(self.g * scale_init)\n+        self.g = self.g * scale_init\n         if hasattr(self.layer, 'bias'):\n-            self.layer.bias.assign(-m_init * scale_init)\n+            self.layer.bias = -m_init * scale_init\n         self.layer.activation = existing_activation\n \n     def get_config(self):\n", "issue": "WeightNormalization fails after update\nYesterday I cut a 0.3 release on pypi test, but as a sanity check on the weightnorm changes #208 I ran a quick regression check:\r\n\r\nhttps://colab.research.google.com/drive/1RQlVR9X7vj8q3W3sJC6YA3Cq45XTXrau\r\n\r\nThe new layer is failing with loss of nan. Looking at the test coverage for this layer I'm struggling how to best catch this without running a full model. The losses for our toy test cases do not return nan so perhaps we need a true regression test like mnist?\r\n\r\nEDIT -- Fix colab link\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\nfrom tensorflow_addons.utils import keras_utils\n\n\n@keras_utils.register_keras_custom_object\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._initialized = False\n        self._track_trackable(layer, name='layer')\n\n    def build(self, input_shape):\n        \"\"\"Build `Layer`\"\"\"\n        input_shape = tf.TensorShape(input_shape).as_list()\n        self.input_spec = tf.keras.layers.InputSpec(shape=input_shape)\n\n        if not self.layer.built:\n            self.layer.build(input_shape)\n\n            if not hasattr(self.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(self.layer.kernel.shape[-1])\n            self.kernel_norm_axes = list(\n                range(self.layer.kernel.shape.rank - 1))\n\n            self.v = self.layer.kernel\n            self.g = self.add_variable(\n                name=\"g\",\n                shape=(self.layer_depth,),\n                initializer=tf.keras.initializers.get('ones'),\n                dtype=self.layer.kernel.dtype,\n                trainable=True)\n\n        super(WeightNormalization, self).build()\n\n    def call(self, inputs):\n        \"\"\"Call `Layer`\"\"\"\n        if not self._initialized:\n            self._initialize_weights(inputs)\n\n        self._compute_weights()  # Recompute weights for each forward pass\n        output = self.layer(inputs)\n        return output\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 _compute_weights(self):\n        \"\"\"Generate normalized weights.\n\n        This method will update the value of self.layer.kernel with the\n        normalized value, so that the layer is ready for call().\n        \"\"\"\n        with tf.name_scope('compute_weights'):\n            self.layer.kernel = tf.nn.l2_normalize(\n                self.v, axis=self.kernel_norm_axes) * self.g\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        if self.data_init:\n            self._data_dep_init(inputs)\n        else:\n            self._init_norm()\n        self._initialized = True\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            flat = tf.reshape(self.v, [-1, self.layer_depth])\n            self.g.assign(\n                tf.reshape(tf.linalg.norm(flat, axis=0), (self.layer_depth,)))\n\n    def _data_dep_init(self, inputs):\n        \"\"\"Data dependent initialization.\"\"\"\n\n        with tf.name_scope('data_dep_init'):\n            # Generate data dependent init values\n            existing_activation = self.layer.activation\n            self.layer.activation = None\n            x_init = self.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        self.g.assign(self.g * scale_init)\n        if hasattr(self.layer, 'bias'):\n            self.layer.bias.assign(-m_init * scale_init)\n        self.layer.activation = existing_activation\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\nfrom tensorflow_addons.utils import keras_utils\n\n\n@keras_utils.register_keras_custom_object\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._initialized = False\n        self._track_trackable(layer, name='layer')\n\n    def build(self, input_shape):\n        \"\"\"Build `Layer`\"\"\"\n        input_shape = tf.TensorShape(input_shape).as_list()\n        self.input_spec = tf.keras.layers.InputSpec(shape=input_shape)\n\n        if not self.layer.built:\n            self.layer.build(input_shape)\n\n            if not hasattr(self.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(self.layer.kernel.shape[-1])\n            self.kernel_norm_axes = list(\n                range(self.layer.kernel.shape.rank - 1))\n\n            self.v = self.layer.kernel\n            self.g = self.add_variable(\n                name=\"g\",\n                shape=(self.layer_depth,),\n                initializer=tf.keras.initializers.get('ones'),\n                dtype=self.layer.kernel.dtype,\n                trainable=True)\n\n        super(WeightNormalization, self).build()\n\n    @tf.function\n    def call(self, inputs):\n        \"\"\"Call `Layer`\"\"\"\n        if not self._initialized:\n            self._initialize_weights(inputs)\n\n        self._compute_weights()  # Recompute weights for each forward pass\n        output = self.layer(inputs)\n        return output\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 _compute_weights(self):\n        \"\"\"Generate normalized weights.\n\n        This method will update the value of self.layer.kernel with the\n        normalized value, so that the layer is ready for call().\n        \"\"\"\n        with tf.name_scope('compute_weights'):\n            self.layer.kernel = tf.nn.l2_normalize(\n                self.v, axis=self.kernel_norm_axes) * self.g\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        if self.data_init:\n            self._data_dep_init(inputs)\n        else:\n            self._init_norm()\n        self._initialized = True\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            flat = tf.reshape(self.v, [-1, self.layer_depth])\n            self.g.assign(\n                tf.reshape(tf.linalg.norm(flat, axis=0), (self.layer_depth,)))\n\n    def _data_dep_init(self, inputs):\n        \"\"\"Data dependent initialization.\"\"\"\n\n        with tf.name_scope('data_dep_init'):\n            # Generate data dependent init values\n            existing_activation = self.layer.activation\n            self.layer.activation = None\n            x_init = self.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        self.g = self.g * scale_init\n        if hasattr(self.layer, 'bias'):\n            self.layer.bias = -m_init * scale_init\n        self.layer.activation = existing_activation\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"}]}