We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Markdown parsing error in NovoGrad documentation page
The documentation page of NovoGrad has a broken Markdown link shown below:

![novograd](https://user-images.githubusercontent.com/6107926/89725432-69760a00-da4a-11ea-9c8c-1495ecf9ec69.png)

I suspect this is because only the first line of the docstring is used for the first line of the documentation page. In the docstring (shown below), the link spans multiple lines.

https://github.com/tensorflow/addons/blob/5f746971d0d9491716f2f13206299a2c45941b0c/tensorflow_addons/optimizers/novograd.py#L26-L28

The simplest solution would be to edit the docstring so that there is a one-line summary for the first line with a paper link after the newline. One example would be [LAMB's documentation page](https://www.tensorflow.org/addons/api_docs/python/tfa/optimizers/LAMB). I will create a PR with this solution.


--- 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/optimizers/novograd.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 """NovoGrad for TensorFlow."""
16 
17 import tensorflow as tf
18 from tensorflow_addons.utils.types import FloatTensorLike
19 
20 from typing import Union, Callable
21 from typeguard import typechecked
22 
23 
24 @tf.keras.utils.register_keras_serializable(package="Addons")
25 class NovoGrad(tf.keras.optimizers.Optimizer):
26     """The NovoGrad Optimizer was first proposed in [Stochastic Gradient
27     Methods with Layerwise Adaptvie Moments for training of Deep
28     Networks](https://arxiv.org/pdf/1905.11286.pdf)
29 
30     NovoGrad is a first-order SGD-based algorithm, which computes second
31     moments per layer instead of per weight as in Adam. Compared to Adam,
32     NovoGrad takes less memory, and has been found to be more numerically
33     stable. More specifically we compute (for more information on the
34     computation please refer to this
35     [link](https://nvidia.github.io/OpenSeq2Seq/html/optimizers.html):
36 
37     Second order moment = exponential moving average of Layer-wise square
38     of grads:
39         v_t <-- beta_2 * v_{t-1} + (1-beta_2) * (g_t)^2
40     First order moment in one of four modes:
41         1. moment of grads normalized by v_t:
42             m_t <- beta_1 * m_{t-1} + [ g_t / (sqrt(v_t)+epsilon)]
43         2. moment similar to Adam: exponential moving average of grads
44         normalized by v_t (set grad_averaging = True to use this):
45             m_t <- beta_1 * m_{t-1} +
46                    [(1 - beta_1) * (g_t / (sqrt(v_t) + epsilon))]
47         3. weight decay adds a w_d term after grads are rescaled by
48         1/sqrt(v_t) (set weight_decay > 0 to use this0:
49             m_t <- beta_1 * m_{t-1} +
50                    [(g_t / (sqrt(v_t) + epsilon)) + (w_d * w_{t-1})]
51         4. weight decay + exponential moving average from Adam:
52             m_t <- beta_1 * m_{t-1} +
53                    [(1 - beta_1) * ((g_t / (sqrt(v_t + epsilon)) +
54                    (w_d * w_{t-1}))]
55     Weight update:
56         w_t <- w_{t-1} - lr_t * m_t
57 
58     Example of usage:
59     ```python
60     opt = tfa.optimizers.NovoGrad(
61         lr=1e-3,
62         beta_1=0.9,
63         beta_2=0.999,
64         weight_decay=0.001,
65         grad_averaging=False
66     )
67     ```
68     """
69 
70     @typechecked
71     def __init__(
72         self,
73         learning_rate: Union[FloatTensorLike, Callable] = 0.001,
74         beta_1: FloatTensorLike = 0.9,
75         beta_2: FloatTensorLike = 0.999,
76         epsilon: FloatTensorLike = 1e-7,
77         weight_decay: FloatTensorLike = 0.0,
78         grad_averaging: bool = False,
79         amsgrad: bool = False,
80         name: str = "NovoGrad",
81         **kwargs
82     ):
83         r"""Construct a new NovoGrad optimizer.
84 
85         Args:
86             learning_rate: A `Tensor` or a floating point value. or a schedule
87                 that is a `tf.keras.optimizers.schedules.LearningRateSchedule`
88                 The learning rate.
89             beta_1: A float value or a constant float tensor.
90                 The exponential decay rate for the 1st moment estimates.
91             beta_2: A float value or a constant float tensor.
92                 The exponential decay rate for the 2nd moment estimates.
93             epsilon: A small constant for numerical stability.
94             weight_decay: A floating point value. Weight decay for each param.
95             grad_averaging: determines whether to use Adam style exponential
96                 moving averaging for the first order moments.
97             **kwargs: keyword arguments. Allowed to be {`clipnorm`,
98                 `clipvalue`, `lr`, `decay`}. `clipnorm` is clip gradients
99                 by norm; `clipvalue` is clip gradients by value, `decay` is
100                 included for backward compatibility to allow time inverse
101                 decay of learning rate. `lr` is included for backward
102                 compatibility, recommended to use `learning_rate` instead.
103         """
104         super().__init__(name, **kwargs)
105         if weight_decay < 0.0:
106             raise ValueError("Weight decay rate cannot be negative")
107         self._set_hyper("learning_rate", kwargs.get("lr", learning_rate))
108         self._set_hyper("decay", self._initial_decay)
109         self._set_hyper("beta_1", beta_1)
110         self._set_hyper("beta_2", beta_2)
111         self._set_hyper("weight_decay", weight_decay)
112         self._set_hyper("grad_averaging", grad_averaging)
113         self.amsgrad = amsgrad
114         self.epsilon = epsilon or tf.keras.backend.epsilon()
115 
116     def _create_slots(self, var_list):
117         # Create slots for the first and second moments.
118         # Separate for-loops to respect the ordering of slot variables from v1.
119         for var in var_list:
120             self.add_slot(var=var, slot_name="m", initializer="zeros")
121         for var in var_list:
122             self.add_slot(
123                 var=var, slot_name="v", initializer=tf.zeros(shape=[], dtype=var.dtype)
124             )
125         if self.amsgrad:
126             for var in var_list:
127                 self.add_slot(var, "vhat")
128 
129     def _prepare_local(self, var_device, var_dtype, apply_state):
130         super()._prepare_local(var_device, var_dtype, apply_state)
131         beta_1_t = tf.identity(self._get_hyper("beta_1", var_dtype))
132         beta_2_t = tf.identity(self._get_hyper("beta_2", var_dtype))
133         apply_state[(var_device, var_dtype)].update(
134             dict(
135                 epsilon=tf.convert_to_tensor(self.epsilon, var_dtype),
136                 beta_1_t=beta_1_t,
137                 beta_2_t=beta_2_t,
138                 one_minus_beta_2_t=1 - beta_2_t,
139                 one_minus_beta_1_t=1 - beta_1_t,
140             )
141         )
142 
143     def set_weights(self, weights):
144         params = self.weights
145         # If the weights are generated by Keras V1 optimizer, it includes vhats
146         # even without amsgrad, i.e, V1 optimizer has 3x + 1 variables, while V2
147         # optimizer has 2x + 1 variables. Filter vhats out for compatibility.
148         num_vars = int((len(params) - 1) / 2)
149         if len(weights) == 3 * num_vars + 1:
150             weights = weights[: len(params)]
151         super().set_weights(weights)
152 
153     def _resource_apply_dense(self, grad, var, apply_state=None):
154         var_device, var_dtype = var.device, var.dtype.base_dtype
155         coefficients = (apply_state or {}).get(
156             (var_device, var_dtype)
157         ) or self._fallback_apply_state(var_device, var_dtype)
158         weight_decay = self._get_hyper("weight_decay")
159         grad_averaging = self._get_hyper("grad_averaging")
160 
161         v = self.get_slot(var, "v")
162         g_2 = tf.reduce_sum(tf.square(tf.cast(grad, tf.float32)))
163         v_t = tf.cond(
164             tf.equal(self.iterations, 0),
165             lambda: g_2,
166             lambda: v * coefficients["beta_2_t"]
167             + g_2 * coefficients["one_minus_beta_2_t"],
168         )
169         v_t = v.assign(v_t, use_locking=self._use_locking)
170 
171         if self.amsgrad:
172             vhat = self.get_slot(var, "vhat")
173             vhat_t = vhat.assign(tf.maximum(vhat, v_t), use_locking=self._use_locking)
174             grad = grad / (tf.sqrt(vhat_t) + self.epsilon)
175         else:
176             grad = grad / (tf.sqrt(v_t) + self.epsilon)
177         grad = tf.cond(
178             tf.greater(weight_decay, 0), lambda: grad + weight_decay * var, lambda: grad
179         )
180         grad = tf.cond(
181             tf.logical_and(grad_averaging, tf.not_equal(self.iterations, 0)),
182             lambda: grad * coefficients["one_minus_beta_1_t"],
183             lambda: grad,
184         )
185         m = self.get_slot(var, "m")
186         return tf.raw_ops.ResourceApplyKerasMomentum(
187             var=var.handle,
188             accum=m.handle,
189             lr=coefficients["lr_t"],
190             grad=grad,
191             momentum=coefficients["beta_1_t"],
192             use_locking=self._use_locking,
193             use_nesterov=False,
194         )
195 
196     def _resource_apply_sparse(self, grad, var, indices, apply_state=None):
197         var_device, var_dtype = var.device, var.dtype.base_dtype
198         coefficients = (apply_state or {}).get(
199             (var_device, var_dtype)
200         ) or self._fallback_apply_state(var_device, var_dtype)
201         weight_decay = self._get_hyper("weight_decay")
202         grad_averaging = self._get_hyper("grad_averaging")
203 
204         v = self.get_slot(var, "v")
205         g_2 = tf.reduce_sum(tf.square(tf.cast(grad, tf.float32)))
206         # v is just a scalar and does not need to involve sparse tensors.
207         v_t = tf.cond(
208             tf.equal(self.iterations, 0),
209             lambda: g_2,
210             lambda: v * coefficients["beta_2_t"]
211             + g_2 * coefficients["one_minus_beta_2_t"],
212         )
213         v_t = v.assign(v_t, use_locking=self._use_locking)
214 
215         if self.amsgrad:
216             vhat = self.get_slot(var, "vhat")
217             vhat_t = vhat.assign(tf.maximum(vhat, v_t), use_locking=self._use_locking)
218             grad = grad / (tf.sqrt(vhat_t) + self.epsilon)
219         else:
220             grad = grad / (tf.sqrt(v_t) + self.epsilon)
221         grad = tf.cond(
222             tf.greater(weight_decay, 0),
223             lambda: grad + weight_decay * tf.gather(var, indices),
224             lambda: grad,
225         )
226         grad = tf.cond(
227             tf.logical_and(grad_averaging, tf.not_equal(self.iterations, 0)),
228             lambda: grad * coefficients["one_minus_beta_1_t"],
229             lambda: grad,
230         )
231         m = self.get_slot(var, "m")
232         return tf.raw_ops.ResourceSparseApplyKerasMomentum(
233             var=var.handle,
234             accum=m.handle,
235             lr=coefficients["lr_t"],
236             grad=grad,
237             indices=indices,
238             momentum=coefficients["beta_1_t"],
239             use_locking=self._use_locking,
240             use_nesterov=False,
241         )
242 
243     def get_config(self):
244         config = super().get_config()
245         config.update(
246             {
247                 "learning_rate": self._serialize_hyperparameter("learning_rate"),
248                 "beta_1": self._serialize_hyperparameter("beta_1"),
249                 "beta_2": self._serialize_hyperparameter("beta_2"),
250                 "epsilon": self.epsilon,
251                 "weight_decay": self._serialize_hyperparameter("weight_decay"),
252                 "grad_averaging": self._serialize_hyperparameter("grad_averaging"),
253             }
254         )
255         return config
256
```
--- 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/optimizers/novograd.py b/tensorflow_addons/optimizers/novograd.py
--- a/tensorflow_addons/optimizers/novograd.py
+++ b/tensorflow_addons/optimizers/novograd.py
@@ -23,16 +23,16 @@
 
 @tf.keras.utils.register_keras_serializable(package="Addons")
 class NovoGrad(tf.keras.optimizers.Optimizer):
-    """The NovoGrad Optimizer was first proposed in [Stochastic Gradient
-    Methods with Layerwise Adaptvie Moments for training of Deep
-    Networks](https://arxiv.org/pdf/1905.11286.pdf)
-
-    NovoGrad is a first-order SGD-based algorithm, which computes second
-    moments per layer instead of per weight as in Adam. Compared to Adam,
-    NovoGrad takes less memory, and has been found to be more numerically
-    stable. More specifically we compute (for more information on the
-    computation please refer to this
-    [link](https://nvidia.github.io/OpenSeq2Seq/html/optimizers.html):
+    """Optimizer that implements NovoGrad.
+
+    The NovoGrad Optimizer was first proposed in [Stochastic Gradient
+    Methods with Layerwise Adaptive Moments for training of Deep
+    Networks](https://arxiv.org/pdf/1905.11286.pdf) NovoGrad is a
+    first-order SGD-based algorithm, which computes second moments per
+    layer instead of per weight as in Adam. Compared to Adam, NovoGrad
+    takes less memory, and has been found to be more numerically stable.
+    (For more information on the computation please refer to this
+    [link](https://nvidia.github.io/OpenSeq2Seq/html/optimizers.html))
 
     Second order moment = exponential moving average of Layer-wise square
     of grads:

{"golden_diff": "diff --git a/tensorflow_addons/optimizers/novograd.py b/tensorflow_addons/optimizers/novograd.py\n--- a/tensorflow_addons/optimizers/novograd.py\n+++ b/tensorflow_addons/optimizers/novograd.py\n@@ -23,16 +23,16 @@\n \n @tf.keras.utils.register_keras_serializable(package=\"Addons\")\n class NovoGrad(tf.keras.optimizers.Optimizer):\n-    \"\"\"The NovoGrad Optimizer was first proposed in [Stochastic Gradient\n-    Methods with Layerwise Adaptvie Moments for training of Deep\n-    Networks](https://arxiv.org/pdf/1905.11286.pdf)\n-\n-    NovoGrad is a first-order SGD-based algorithm, which computes second\n-    moments per layer instead of per weight as in Adam. Compared to Adam,\n-    NovoGrad takes less memory, and has been found to be more numerically\n-    stable. More specifically we compute (for more information on the\n-    computation please refer to this\n-    [link](https://nvidia.github.io/OpenSeq2Seq/html/optimizers.html):\n+    \"\"\"Optimizer that implements NovoGrad.\n+\n+    The NovoGrad Optimizer was first proposed in [Stochastic Gradient\n+    Methods with Layerwise Adaptive Moments for training of Deep\n+    Networks](https://arxiv.org/pdf/1905.11286.pdf) NovoGrad is a\n+    first-order SGD-based algorithm, which computes second moments per\n+    layer instead of per weight as in Adam. Compared to Adam, NovoGrad\n+    takes less memory, and has been found to be more numerically stable.\n+    (For more information on the computation please refer to this\n+    [link](https://nvidia.github.io/OpenSeq2Seq/html/optimizers.html))\n \n     Second order moment = exponential moving average of Layer-wise square\n     of grads:\n", "issue": "Markdown parsing error in NovoGrad documentation page\nThe documentation page of NovoGrad has a broken Markdown link shown below:\r\n\r\n![novograd](https://user-images.githubusercontent.com/6107926/89725432-69760a00-da4a-11ea-9c8c-1495ecf9ec69.png)\r\n\r\nI suspect this is because only the first line of the docstring is used for the first line of the documentation page. In the docstring (shown below), the link spans multiple lines.\r\n\r\nhttps://github.com/tensorflow/addons/blob/5f746971d0d9491716f2f13206299a2c45941b0c/tensorflow_addons/optimizers/novograd.py#L26-L28\r\n\r\nThe simplest solution would be to edit the docstring so that there is a one-line summary for the first line with a paper link after the newline. One example would be [LAMB's documentation page](https://www.tensorflow.org/addons/api_docs/python/tfa/optimizers/LAMB). I will create a PR with this solution.\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# ==============================================================================\n\"\"\"NovoGrad for TensorFlow.\"\"\"\n\nimport tensorflow as tf\nfrom tensorflow_addons.utils.types import FloatTensorLike\n\nfrom typing import Union, Callable\nfrom typeguard import typechecked\n\n\n@tf.keras.utils.register_keras_serializable(package=\"Addons\")\nclass NovoGrad(tf.keras.optimizers.Optimizer):\n    \"\"\"The NovoGrad Optimizer was first proposed in [Stochastic Gradient\n    Methods with Layerwise Adaptvie Moments for training of Deep\n    Networks](https://arxiv.org/pdf/1905.11286.pdf)\n\n    NovoGrad is a first-order SGD-based algorithm, which computes second\n    moments per layer instead of per weight as in Adam. Compared to Adam,\n    NovoGrad takes less memory, and has been found to be more numerically\n    stable. More specifically we compute (for more information on the\n    computation please refer to this\n    [link](https://nvidia.github.io/OpenSeq2Seq/html/optimizers.html):\n\n    Second order moment = exponential moving average of Layer-wise square\n    of grads:\n        v_t <-- beta_2 * v_{t-1} + (1-beta_2) * (g_t)^2\n    First order moment in one of four modes:\n        1. moment of grads normalized by v_t:\n            m_t <- beta_1 * m_{t-1} + [ g_t / (sqrt(v_t)+epsilon)]\n        2. moment similar to Adam: exponential moving average of grads\n        normalized by v_t (set grad_averaging = True to use this):\n            m_t <- beta_1 * m_{t-1} +\n                   [(1 - beta_1) * (g_t / (sqrt(v_t) + epsilon))]\n        3. weight decay adds a w_d term after grads are rescaled by\n        1/sqrt(v_t) (set weight_decay > 0 to use this0:\n            m_t <- beta_1 * m_{t-1} +\n                   [(g_t / (sqrt(v_t) + epsilon)) + (w_d * w_{t-1})]\n        4. weight decay + exponential moving average from Adam:\n            m_t <- beta_1 * m_{t-1} +\n                   [(1 - beta_1) * ((g_t / (sqrt(v_t + epsilon)) +\n                   (w_d * w_{t-1}))]\n    Weight update:\n        w_t <- w_{t-1} - lr_t * m_t\n\n    Example of usage:\n    ```python\n    opt = tfa.optimizers.NovoGrad(\n        lr=1e-3,\n        beta_1=0.9,\n        beta_2=0.999,\n        weight_decay=0.001,\n        grad_averaging=False\n    )\n    ```\n    \"\"\"\n\n    @typechecked\n    def __init__(\n        self,\n        learning_rate: Union[FloatTensorLike, Callable] = 0.001,\n        beta_1: FloatTensorLike = 0.9,\n        beta_2: FloatTensorLike = 0.999,\n        epsilon: FloatTensorLike = 1e-7,\n        weight_decay: FloatTensorLike = 0.0,\n        grad_averaging: bool = False,\n        amsgrad: bool = False,\n        name: str = \"NovoGrad\",\n        **kwargs\n    ):\n        r\"\"\"Construct a new NovoGrad optimizer.\n\n        Args:\n            learning_rate: A `Tensor` or a floating point value. or a schedule\n                that is a `tf.keras.optimizers.schedules.LearningRateSchedule`\n                The learning rate.\n            beta_1: A float value or a constant float tensor.\n                The exponential decay rate for the 1st moment estimates.\n            beta_2: A float value or a constant float tensor.\n                The exponential decay rate for the 2nd moment estimates.\n            epsilon: A small constant for numerical stability.\n            weight_decay: A floating point value. Weight decay for each param.\n            grad_averaging: determines whether to use Adam style exponential\n                moving averaging for the first order moments.\n            **kwargs: keyword arguments. Allowed to be {`clipnorm`,\n                `clipvalue`, `lr`, `decay`}. `clipnorm` is clip gradients\n                by norm; `clipvalue` is clip gradients by value, `decay` is\n                included for backward compatibility to allow time inverse\n                decay of learning rate. `lr` is included for backward\n                compatibility, recommended to use `learning_rate` instead.\n        \"\"\"\n        super().__init__(name, **kwargs)\n        if weight_decay < 0.0:\n            raise ValueError(\"Weight decay rate cannot be negative\")\n        self._set_hyper(\"learning_rate\", kwargs.get(\"lr\", learning_rate))\n        self._set_hyper(\"decay\", self._initial_decay)\n        self._set_hyper(\"beta_1\", beta_1)\n        self._set_hyper(\"beta_2\", beta_2)\n        self._set_hyper(\"weight_decay\", weight_decay)\n        self._set_hyper(\"grad_averaging\", grad_averaging)\n        self.amsgrad = amsgrad\n        self.epsilon = epsilon or tf.keras.backend.epsilon()\n\n    def _create_slots(self, var_list):\n        # Create slots for the first and second moments.\n        # Separate for-loops to respect the ordering of slot variables from v1.\n        for var in var_list:\n            self.add_slot(var=var, slot_name=\"m\", initializer=\"zeros\")\n        for var in var_list:\n            self.add_slot(\n                var=var, slot_name=\"v\", initializer=tf.zeros(shape=[], dtype=var.dtype)\n            )\n        if self.amsgrad:\n            for var in var_list:\n                self.add_slot(var, \"vhat\")\n\n    def _prepare_local(self, var_device, var_dtype, apply_state):\n        super()._prepare_local(var_device, var_dtype, apply_state)\n        beta_1_t = tf.identity(self._get_hyper(\"beta_1\", var_dtype))\n        beta_2_t = tf.identity(self._get_hyper(\"beta_2\", var_dtype))\n        apply_state[(var_device, var_dtype)].update(\n            dict(\n                epsilon=tf.convert_to_tensor(self.epsilon, var_dtype),\n                beta_1_t=beta_1_t,\n                beta_2_t=beta_2_t,\n                one_minus_beta_2_t=1 - beta_2_t,\n                one_minus_beta_1_t=1 - beta_1_t,\n            )\n        )\n\n    def set_weights(self, weights):\n        params = self.weights\n        # If the weights are generated by Keras V1 optimizer, it includes vhats\n        # even without amsgrad, i.e, V1 optimizer has 3x + 1 variables, while V2\n        # optimizer has 2x + 1 variables. Filter vhats out for compatibility.\n        num_vars = int((len(params) - 1) / 2)\n        if len(weights) == 3 * num_vars + 1:\n            weights = weights[: len(params)]\n        super().set_weights(weights)\n\n    def _resource_apply_dense(self, grad, var, apply_state=None):\n        var_device, var_dtype = var.device, var.dtype.base_dtype\n        coefficients = (apply_state or {}).get(\n            (var_device, var_dtype)\n        ) or self._fallback_apply_state(var_device, var_dtype)\n        weight_decay = self._get_hyper(\"weight_decay\")\n        grad_averaging = self._get_hyper(\"grad_averaging\")\n\n        v = self.get_slot(var, \"v\")\n        g_2 = tf.reduce_sum(tf.square(tf.cast(grad, tf.float32)))\n        v_t = tf.cond(\n            tf.equal(self.iterations, 0),\n            lambda: g_2,\n            lambda: v * coefficients[\"beta_2_t\"]\n            + g_2 * coefficients[\"one_minus_beta_2_t\"],\n        )\n        v_t = v.assign(v_t, use_locking=self._use_locking)\n\n        if self.amsgrad:\n            vhat = self.get_slot(var, \"vhat\")\n            vhat_t = vhat.assign(tf.maximum(vhat, v_t), use_locking=self._use_locking)\n            grad = grad / (tf.sqrt(vhat_t) + self.epsilon)\n        else:\n            grad = grad / (tf.sqrt(v_t) + self.epsilon)\n        grad = tf.cond(\n            tf.greater(weight_decay, 0), lambda: grad + weight_decay * var, lambda: grad\n        )\n        grad = tf.cond(\n            tf.logical_and(grad_averaging, tf.not_equal(self.iterations, 0)),\n            lambda: grad * coefficients[\"one_minus_beta_1_t\"],\n            lambda: grad,\n        )\n        m = self.get_slot(var, \"m\")\n        return tf.raw_ops.ResourceApplyKerasMomentum(\n            var=var.handle,\n            accum=m.handle,\n            lr=coefficients[\"lr_t\"],\n            grad=grad,\n            momentum=coefficients[\"beta_1_t\"],\n            use_locking=self._use_locking,\n            use_nesterov=False,\n        )\n\n    def _resource_apply_sparse(self, grad, var, indices, apply_state=None):\n        var_device, var_dtype = var.device, var.dtype.base_dtype\n        coefficients = (apply_state or {}).get(\n            (var_device, var_dtype)\n        ) or self._fallback_apply_state(var_device, var_dtype)\n        weight_decay = self._get_hyper(\"weight_decay\")\n        grad_averaging = self._get_hyper(\"grad_averaging\")\n\n        v = self.get_slot(var, \"v\")\n        g_2 = tf.reduce_sum(tf.square(tf.cast(grad, tf.float32)))\n        # v is just a scalar and does not need to involve sparse tensors.\n        v_t = tf.cond(\n            tf.equal(self.iterations, 0),\n            lambda: g_2,\n            lambda: v * coefficients[\"beta_2_t\"]\n            + g_2 * coefficients[\"one_minus_beta_2_t\"],\n        )\n        v_t = v.assign(v_t, use_locking=self._use_locking)\n\n        if self.amsgrad:\n            vhat = self.get_slot(var, \"vhat\")\n            vhat_t = vhat.assign(tf.maximum(vhat, v_t), use_locking=self._use_locking)\n            grad = grad / (tf.sqrt(vhat_t) + self.epsilon)\n        else:\n            grad = grad / (tf.sqrt(v_t) + self.epsilon)\n        grad = tf.cond(\n            tf.greater(weight_decay, 0),\n            lambda: grad + weight_decay * tf.gather(var, indices),\n            lambda: grad,\n        )\n        grad = tf.cond(\n            tf.logical_and(grad_averaging, tf.not_equal(self.iterations, 0)),\n            lambda: grad * coefficients[\"one_minus_beta_1_t\"],\n            lambda: grad,\n        )\n        m = self.get_slot(var, \"m\")\n        return tf.raw_ops.ResourceSparseApplyKerasMomentum(\n            var=var.handle,\n            accum=m.handle,\n            lr=coefficients[\"lr_t\"],\n            grad=grad,\n            indices=indices,\n            momentum=coefficients[\"beta_1_t\"],\n            use_locking=self._use_locking,\n            use_nesterov=False,\n        )\n\n    def get_config(self):\n        config = super().get_config()\n        config.update(\n            {\n                \"learning_rate\": self._serialize_hyperparameter(\"learning_rate\"),\n                \"beta_1\": self._serialize_hyperparameter(\"beta_1\"),\n                \"beta_2\": self._serialize_hyperparameter(\"beta_2\"),\n                \"epsilon\": self.epsilon,\n                \"weight_decay\": self._serialize_hyperparameter(\"weight_decay\"),\n                \"grad_averaging\": self._serialize_hyperparameter(\"grad_averaging\"),\n            }\n        )\n        return config\n", "path": "tensorflow_addons/optimizers/novograd.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# ==============================================================================\n\"\"\"NovoGrad for TensorFlow.\"\"\"\n\nimport tensorflow as tf\nfrom tensorflow_addons.utils.types import FloatTensorLike\n\nfrom typing import Union, Callable\nfrom typeguard import typechecked\n\n\n@tf.keras.utils.register_keras_serializable(package=\"Addons\")\nclass NovoGrad(tf.keras.optimizers.Optimizer):\n    \"\"\"Optimizer that implements NovoGrad.\n\n    The NovoGrad Optimizer was first proposed in [Stochastic Gradient\n    Methods with Layerwise Adaptive Moments for training of Deep\n    Networks](https://arxiv.org/pdf/1905.11286.pdf) NovoGrad is a\n    first-order SGD-based algorithm, which computes second moments per\n    layer instead of per weight as in Adam. Compared to Adam, NovoGrad\n    takes less memory, and has been found to be more numerically stable.\n    (For more information on the computation please refer to this\n    [link](https://nvidia.github.io/OpenSeq2Seq/html/optimizers.html))\n\n    Second order moment = exponential moving average of Layer-wise square\n    of grads:\n        v_t <-- beta_2 * v_{t-1} + (1-beta_2) * (g_t)^2\n    First order moment in one of four modes:\n        1. moment of grads normalized by v_t:\n            m_t <- beta_1 * m_{t-1} + [ g_t / (sqrt(v_t)+epsilon)]\n        2. moment similar to Adam: exponential moving average of grads\n        normalized by v_t (set grad_averaging = True to use this):\n            m_t <- beta_1 * m_{t-1} +\n                   [(1 - beta_1) * (g_t / (sqrt(v_t) + epsilon))]\n        3. weight decay adds a w_d term after grads are rescaled by\n        1/sqrt(v_t) (set weight_decay > 0 to use this0:\n            m_t <- beta_1 * m_{t-1} +\n                   [(g_t / (sqrt(v_t) + epsilon)) + (w_d * w_{t-1})]\n        4. weight decay + exponential moving average from Adam:\n            m_t <- beta_1 * m_{t-1} +\n                   [(1 - beta_1) * ((g_t / (sqrt(v_t + epsilon)) +\n                   (w_d * w_{t-1}))]\n    Weight update:\n        w_t <- w_{t-1} - lr_t * m_t\n\n    Example of usage:\n    ```python\n    opt = tfa.optimizers.NovoGrad(\n        lr=1e-3,\n        beta_1=0.9,\n        beta_2=0.999,\n        weight_decay=0.001,\n        grad_averaging=False\n    )\n    ```\n    \"\"\"\n\n    @typechecked\n    def __init__(\n        self,\n        learning_rate: Union[FloatTensorLike, Callable] = 0.001,\n        beta_1: FloatTensorLike = 0.9,\n        beta_2: FloatTensorLike = 0.999,\n        epsilon: FloatTensorLike = 1e-7,\n        weight_decay: FloatTensorLike = 0.0,\n        grad_averaging: bool = False,\n        amsgrad: bool = False,\n        name: str = \"NovoGrad\",\n        **kwargs\n    ):\n        r\"\"\"Construct a new NovoGrad optimizer.\n\n        Args:\n            learning_rate: A `Tensor` or a floating point value. or a schedule\n                that is a `tf.keras.optimizers.schedules.LearningRateSchedule`\n                The learning rate.\n            beta_1: A float value or a constant float tensor.\n                The exponential decay rate for the 1st moment estimates.\n            beta_2: A float value or a constant float tensor.\n                The exponential decay rate for the 2nd moment estimates.\n            epsilon: A small constant for numerical stability.\n            weight_decay: A floating point value. Weight decay for each param.\n            grad_averaging: determines whether to use Adam style exponential\n                moving averaging for the first order moments.\n            **kwargs: keyword arguments. Allowed to be {`clipnorm`,\n                `clipvalue`, `lr`, `decay`}. `clipnorm` is clip gradients\n                by norm; `clipvalue` is clip gradients by value, `decay` is\n                included for backward compatibility to allow time inverse\n                decay of learning rate. `lr` is included for backward\n                compatibility, recommended to use `learning_rate` instead.\n        \"\"\"\n        super().__init__(name, **kwargs)\n        if weight_decay < 0.0:\n            raise ValueError(\"Weight decay rate cannot be negative\")\n        self._set_hyper(\"learning_rate\", kwargs.get(\"lr\", learning_rate))\n        self._set_hyper(\"decay\", self._initial_decay)\n        self._set_hyper(\"beta_1\", beta_1)\n        self._set_hyper(\"beta_2\", beta_2)\n        self._set_hyper(\"weight_decay\", weight_decay)\n        self._set_hyper(\"grad_averaging\", grad_averaging)\n        self.amsgrad = amsgrad\n        self.epsilon = epsilon or tf.keras.backend.epsilon()\n\n    def _create_slots(self, var_list):\n        # Create slots for the first and second moments.\n        # Separate for-loops to respect the ordering of slot variables from v1.\n        for var in var_list:\n            self.add_slot(var=var, slot_name=\"m\", initializer=\"zeros\")\n        for var in var_list:\n            self.add_slot(\n                var=var, slot_name=\"v\", initializer=tf.zeros(shape=[], dtype=var.dtype)\n            )\n        if self.amsgrad:\n            for var in var_list:\n                self.add_slot(var, \"vhat\")\n\n    def _prepare_local(self, var_device, var_dtype, apply_state):\n        super()._prepare_local(var_device, var_dtype, apply_state)\n        beta_1_t = tf.identity(self._get_hyper(\"beta_1\", var_dtype))\n        beta_2_t = tf.identity(self._get_hyper(\"beta_2\", var_dtype))\n        apply_state[(var_device, var_dtype)].update(\n            dict(\n                epsilon=tf.convert_to_tensor(self.epsilon, var_dtype),\n                beta_1_t=beta_1_t,\n                beta_2_t=beta_2_t,\n                one_minus_beta_2_t=1 - beta_2_t,\n                one_minus_beta_1_t=1 - beta_1_t,\n            )\n        )\n\n    def set_weights(self, weights):\n        params = self.weights\n        # If the weights are generated by Keras V1 optimizer, it includes vhats\n        # even without amsgrad, i.e, V1 optimizer has 3x + 1 variables, while V2\n        # optimizer has 2x + 1 variables. Filter vhats out for compatibility.\n        num_vars = int((len(params) - 1) / 2)\n        if len(weights) == 3 * num_vars + 1:\n            weights = weights[: len(params)]\n        super().set_weights(weights)\n\n    def _resource_apply_dense(self, grad, var, apply_state=None):\n        var_device, var_dtype = var.device, var.dtype.base_dtype\n        coefficients = (apply_state or {}).get(\n            (var_device, var_dtype)\n        ) or self._fallback_apply_state(var_device, var_dtype)\n        weight_decay = self._get_hyper(\"weight_decay\")\n        grad_averaging = self._get_hyper(\"grad_averaging\")\n\n        v = self.get_slot(var, \"v\")\n        g_2 = tf.reduce_sum(tf.square(tf.cast(grad, tf.float32)))\n        v_t = tf.cond(\n            tf.equal(self.iterations, 0),\n            lambda: g_2,\n            lambda: v * coefficients[\"beta_2_t\"]\n            + g_2 * coefficients[\"one_minus_beta_2_t\"],\n        )\n        v_t = v.assign(v_t, use_locking=self._use_locking)\n\n        if self.amsgrad:\n            vhat = self.get_slot(var, \"vhat\")\n            vhat_t = vhat.assign(tf.maximum(vhat, v_t), use_locking=self._use_locking)\n            grad = grad / (tf.sqrt(vhat_t) + self.epsilon)\n        else:\n            grad = grad / (tf.sqrt(v_t) + self.epsilon)\n        grad = tf.cond(\n            tf.greater(weight_decay, 0), lambda: grad + weight_decay * var, lambda: grad\n        )\n        grad = tf.cond(\n            tf.logical_and(grad_averaging, tf.not_equal(self.iterations, 0)),\n            lambda: grad * coefficients[\"one_minus_beta_1_t\"],\n            lambda: grad,\n        )\n        m = self.get_slot(var, \"m\")\n        return tf.raw_ops.ResourceApplyKerasMomentum(\n            var=var.handle,\n            accum=m.handle,\n            lr=coefficients[\"lr_t\"],\n            grad=grad,\n            momentum=coefficients[\"beta_1_t\"],\n            use_locking=self._use_locking,\n            use_nesterov=False,\n        )\n\n    def _resource_apply_sparse(self, grad, var, indices, apply_state=None):\n        var_device, var_dtype = var.device, var.dtype.base_dtype\n        coefficients = (apply_state or {}).get(\n            (var_device, var_dtype)\n        ) or self._fallback_apply_state(var_device, var_dtype)\n        weight_decay = self._get_hyper(\"weight_decay\")\n        grad_averaging = self._get_hyper(\"grad_averaging\")\n\n        v = self.get_slot(var, \"v\")\n        g_2 = tf.reduce_sum(tf.square(tf.cast(grad, tf.float32)))\n        # v is just a scalar and does not need to involve sparse tensors.\n        v_t = tf.cond(\n            tf.equal(self.iterations, 0),\n            lambda: g_2,\n            lambda: v * coefficients[\"beta_2_t\"]\n            + g_2 * coefficients[\"one_minus_beta_2_t\"],\n        )\n        v_t = v.assign(v_t, use_locking=self._use_locking)\n\n        if self.amsgrad:\n            vhat = self.get_slot(var, \"vhat\")\n            vhat_t = vhat.assign(tf.maximum(vhat, v_t), use_locking=self._use_locking)\n            grad = grad / (tf.sqrt(vhat_t) + self.epsilon)\n        else:\n            grad = grad / (tf.sqrt(v_t) + self.epsilon)\n        grad = tf.cond(\n            tf.greater(weight_decay, 0),\n            lambda: grad + weight_decay * tf.gather(var, indices),\n            lambda: grad,\n        )\n        grad = tf.cond(\n            tf.logical_and(grad_averaging, tf.not_equal(self.iterations, 0)),\n            lambda: grad * coefficients[\"one_minus_beta_1_t\"],\n            lambda: grad,\n        )\n        m = self.get_slot(var, \"m\")\n        return tf.raw_ops.ResourceSparseApplyKerasMomentum(\n            var=var.handle,\n            accum=m.handle,\n            lr=coefficients[\"lr_t\"],\n            grad=grad,\n            indices=indices,\n            momentum=coefficients[\"beta_1_t\"],\n            use_locking=self._use_locking,\n            use_nesterov=False,\n        )\n\n    def get_config(self):\n        config = super().get_config()\n        config.update(\n            {\n                \"learning_rate\": self._serialize_hyperparameter(\"learning_rate\"),\n                \"beta_1\": self._serialize_hyperparameter(\"beta_1\"),\n                \"beta_2\": self._serialize_hyperparameter(\"beta_2\"),\n                \"epsilon\": self.epsilon,\n                \"weight_decay\": self._serialize_hyperparameter(\"weight_decay\"),\n                \"grad_averaging\": self._serialize_hyperparameter(\"grad_averaging\"),\n            }\n        )\n        return config\n", "path": "tensorflow_addons/optimizers/novograd.py"}]}