We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Lookahead might divide weights with mixed_float16
**System information**
- OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Ubuntu 18.04
- TensorFlow version and how it was installed (source or binary): 2.3.0-rc2 (source)
- TensorFlow-Addons version and how it was installed (source or binary): Just lookahead extracted from the source (master)
- Python version: 3.6.9
- Is GPU used? (yes/no): yes

When tf.keras.mixed_precision.experimental.set_policy('mixed_float16') is set, using LookAhead divides the layers weights after a few gradient updates (for example 7 in my case).

After investigating, I think the reason is that when mixed_float16 is used with automatic loss scaling, the loss scale is initialized very high. Thus the first few gradient updates do nothing (see https://github.com/tensorflow/tensorflow/blob/582c8d236cb079023657287c318ff26adb239002/tensorflow/python/keras/mixed_precision/loss_scale_optimizer.py#L714 for how updates are skipped).

As a result, LookAhead ends up with its 'slow vars' initialized to 0 instead of the original layer weight. Indeed the vars are initialized to the correct value only on the first iteration during the gradient update (https://github.com/tensorflow/addons/blob/master/tensorflow_addons/optimizers/lookahead.py#L106).

Thus when the gradient update finally does something, the first variable update does average them with the slow variables (initialized to zero). Which would explain the weight division I'm seeing.

I don't have a short code to replicate the bug, but I think the above explanation and links should highlight well enough the issue with the current way LookAhead's slow variables are initialized.

Lookahead might divide weights with mixed_float16
**System information**
- OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Ubuntu 18.04
- TensorFlow version and how it was installed (source or binary): 2.3.0-rc2 (source)
- TensorFlow-Addons version and how it was installed (source or binary): Just lookahead extracted from the source (master)
- Python version: 3.6.9
- Is GPU used? (yes/no): yes

When tf.keras.mixed_precision.experimental.set_policy('mixed_float16') is set, using LookAhead divides the layers weights after a few gradient updates (for example 7 in my case).

After investigating, I think the reason is that when mixed_float16 is used with automatic loss scaling, the loss scale is initialized very high. Thus the first few gradient updates do nothing (see https://github.com/tensorflow/tensorflow/blob/582c8d236cb079023657287c318ff26adb239002/tensorflow/python/keras/mixed_precision/loss_scale_optimizer.py#L714 for how updates are skipped).

As a result, LookAhead ends up with its 'slow vars' initialized to 0 instead of the original layer weight. Indeed the vars are initialized to the correct value only on the first iteration during the gradient update (https://github.com/tensorflow/addons/blob/master/tensorflow_addons/optimizers/lookahead.py#L106).

Thus when the gradient update finally does something, the first variable update does average them with the slow variables (initialized to zero). Which would explain the weight division I'm seeing.

I don't have a short code to replicate the bug, but I think the above explanation and links should highlight well enough the issue with the current way LookAhead's slow variables are initialized.


--- 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/lookahead.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 
16 import tensorflow as tf
17 from tensorflow_addons.utils import types
18 
19 from typeguard import typechecked
20 
21 
22 @tf.keras.utils.register_keras_serializable(package="Addons")
23 class Lookahead(tf.keras.optimizers.Optimizer):
24     """This class allows to extend optimizers with the lookahead mechanism.
25 
26     The mechanism is proposed by Michael R. Zhang et.al in the paper
27     [Lookahead Optimizer: k steps forward, 1 step back]
28     (https://arxiv.org/abs/1907.08610v1). The optimizer iteratively updates two
29     sets of weights: the search directions for weights are chosen by the inner
30     optimizer, while the "slow weights" are updated each `k` steps based on the
31     directions of the "fast weights" and the two sets of weights are
32     synchronized. This method improves the learning stability and lowers the
33     variance of its inner optimizer.
34 
35     Example of usage:
36 
37     ```python
38     opt = tf.keras.optimizers.SGD(learning_rate)
39     opt = tfa.optimizers.Lookahead(opt)
40     ```
41     """
42 
43     @typechecked
44     def __init__(
45         self,
46         optimizer: types.Optimizer,
47         sync_period: int = 6,
48         slow_step_size: types.FloatTensorLike = 0.5,
49         name: str = "Lookahead",
50         **kwargs,
51     ):
52         r"""Wrap optimizer with the lookahead mechanism.
53 
54         Args:
55             optimizer: The original optimizer that will be used to compute
56                 and apply the gradients.
57             sync_period: An integer. The synchronization period of lookahead.
58                 Enable lookahead mechanism by setting it with a positive value.
59             slow_step_size: A floating point value.
60                 The ratio for updating the slow weights.
61             name: Optional name for the operations created when applying
62                 gradients. Defaults to "Lookahead".
63             **kwargs: keyword arguments. Allowed to be {`clipnorm`,
64                 `clipvalue`, `lr`, `decay`}. `clipnorm` is clip gradients
65                 by norm; `clipvalue` is clip gradients by value, `decay` is
66                 included for backward compatibility to allow time inverse
67                 decay of learning rate. `lr` is included for backward
68                 compatibility, recommended to use `learning_rate` instead.
69         """
70         super().__init__(name, **kwargs)
71 
72         if isinstance(optimizer, str):
73             optimizer = tf.keras.optimizers.get(optimizer)
74         if not isinstance(optimizer, tf.keras.optimizers.Optimizer):
75             raise TypeError(
76                 "optimizer is not an object of tf.keras.optimizers.Optimizer"
77             )
78 
79         self._optimizer = optimizer
80         self._set_hyper("sync_period", sync_period)
81         self._set_hyper("slow_step_size", slow_step_size)
82         self._initialized = False
83         self._track_trackable(self._optimizer, "lh_base_optimizer")
84 
85     def _create_slots(self, var_list):
86         self._optimizer._create_slots(
87             var_list=var_list
88         )  # pylint: disable=protected-access
89         for var in var_list:
90             self.add_slot(var, "slow")
91 
92     def _create_hypers(self):
93         self._optimizer._create_hypers()  # pylint: disable=protected-access
94 
95     def _prepare(self, var_list):
96         return self._optimizer._prepare(
97             var_list=var_list
98         )  # pylint: disable=protected-access
99 
100     def apply_gradients(self, grads_and_vars, name=None, **kwargs):
101         self._optimizer._iterations = (
102             self.iterations
103         )  # pylint: disable=protected-access
104         return super().apply_gradients(grads_and_vars, name, **kwargs)
105 
106     def _init_op(self, var):
107         slow_var = self.get_slot(var, "slow")
108         return slow_var.assign(
109             tf.where(
110                 tf.equal(self.iterations, tf.constant(0, dtype=self.iterations.dtype)),
111                 var,
112                 slow_var,
113             ),
114             use_locking=self._use_locking,
115         )
116 
117     def _look_ahead_op(self, var):
118         var_dtype = var.dtype.base_dtype
119         slow_var = self.get_slot(var, "slow")
120         local_step = tf.cast(self.iterations + 1, tf.dtypes.int64)
121         sync_period = self._get_hyper("sync_period", tf.dtypes.int64)
122         slow_step_size = self._get_hyper("slow_step_size", var_dtype)
123         step_back = slow_var + slow_step_size * (var - slow_var)
124         sync_cond = tf.equal(
125             tf.math.floordiv(local_step, sync_period) * sync_period, local_step
126         )
127         with tf.control_dependencies([step_back]):
128             slow_update = slow_var.assign(
129                 tf.where(sync_cond, step_back, slow_var), use_locking=self._use_locking
130             )
131             var_update = var.assign(
132                 tf.where(sync_cond, step_back, var), use_locking=self._use_locking
133             )
134         return tf.group(slow_update, var_update)
135 
136     @property
137     def weights(self):
138         return self._weights + self._optimizer.weights
139 
140     def _resource_apply_dense(self, grad, var):
141         init_op = self._init_op(var)
142         with tf.control_dependencies([init_op]):
143             train_op = self._optimizer._resource_apply_dense(
144                 grad, var
145             )  # pylint: disable=protected-access
146             with tf.control_dependencies([train_op]):
147                 look_ahead_op = self._look_ahead_op(var)
148         return tf.group(init_op, train_op, look_ahead_op)
149 
150     def _resource_apply_sparse(self, grad, var, indices):
151         init_op = self._init_op(var)
152         with tf.control_dependencies([init_op]):
153             train_op = self._optimizer._resource_apply_sparse(  # pylint: disable=protected-access
154                 grad, var, indices
155             )
156             with tf.control_dependencies([train_op]):
157                 look_ahead_op = self._look_ahead_op(var)
158         return tf.group(init_op, train_op, look_ahead_op)
159 
160     def get_config(self):
161         config = {
162             "optimizer": tf.keras.optimizers.serialize(self._optimizer),
163             "sync_period": self._serialize_hyperparameter("sync_period"),
164             "slow_step_size": self._serialize_hyperparameter("slow_step_size"),
165         }
166         base_config = super().get_config()
167         return {**base_config, **config}
168 
169     @property
170     def learning_rate(self):
171         return self._optimizer._get_hyper("learning_rate")
172 
173     @learning_rate.setter
174     def learning_rate(self, learning_rate):
175         self._optimizer._set_hyper("learning_rate", learning_rate)
176 
177     @property
178     def lr(self):
179         return self.learning_rate
180 
181     @lr.setter
182     def lr(self, lr):
183         self.learning_rate = lr
184 
185     @classmethod
186     def from_config(cls, config, custom_objects=None):
187         optimizer = tf.keras.optimizers.deserialize(
188             config.pop("optimizer"), custom_objects=custom_objects
189         )
190         return cls(optimizer, **config)
191
```
--- 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/lookahead.py b/tensorflow_addons/optimizers/lookahead.py
--- a/tensorflow_addons/optimizers/lookahead.py
+++ b/tensorflow_addons/optimizers/lookahead.py
@@ -87,7 +87,7 @@
             var_list=var_list
         )  # pylint: disable=protected-access
         for var in var_list:
-            self.add_slot(var, "slow")
+            self.add_slot(var, "slow", initializer=var)
 
     def _create_hypers(self):
         self._optimizer._create_hypers()  # pylint: disable=protected-access
@@ -103,17 +103,6 @@
         )  # pylint: disable=protected-access
         return super().apply_gradients(grads_and_vars, name, **kwargs)
 
-    def _init_op(self, var):
-        slow_var = self.get_slot(var, "slow")
-        return slow_var.assign(
-            tf.where(
-                tf.equal(self.iterations, tf.constant(0, dtype=self.iterations.dtype)),
-                var,
-                slow_var,
-            ),
-            use_locking=self._use_locking,
-        )
-
     def _look_ahead_op(self, var):
         var_dtype = var.dtype.base_dtype
         slow_var = self.get_slot(var, "slow")
@@ -138,24 +127,22 @@
         return self._weights + self._optimizer.weights
 
     def _resource_apply_dense(self, grad, var):
-        init_op = self._init_op(var)
-        with tf.control_dependencies([init_op]):
-            train_op = self._optimizer._resource_apply_dense(
-                grad, var
-            )  # pylint: disable=protected-access
-            with tf.control_dependencies([train_op]):
-                look_ahead_op = self._look_ahead_op(var)
-        return tf.group(init_op, train_op, look_ahead_op)
+        train_op = self._optimizer._resource_apply_dense(
+            grad, var
+        )  # pylint: disable=protected-access
+        with tf.control_dependencies([train_op]):
+            look_ahead_op = self._look_ahead_op(var)
+        return tf.group(train_op, look_ahead_op)
 
     def _resource_apply_sparse(self, grad, var, indices):
-        init_op = self._init_op(var)
-        with tf.control_dependencies([init_op]):
-            train_op = self._optimizer._resource_apply_sparse(  # pylint: disable=protected-access
+        train_op = (
+            self._optimizer._resource_apply_sparse(  # pylint: disable=protected-access
                 grad, var, indices
             )
-            with tf.control_dependencies([train_op]):
-                look_ahead_op = self._look_ahead_op(var)
-        return tf.group(init_op, train_op, look_ahead_op)
+        )
+        with tf.control_dependencies([train_op]):
+            look_ahead_op = self._look_ahead_op(var)
+        return tf.group(train_op, look_ahead_op)
 
     def get_config(self):
         config = {

{"golden_diff": "diff --git a/tensorflow_addons/optimizers/lookahead.py b/tensorflow_addons/optimizers/lookahead.py\n--- a/tensorflow_addons/optimizers/lookahead.py\n+++ b/tensorflow_addons/optimizers/lookahead.py\n@@ -87,7 +87,7 @@\n             var_list=var_list\n         )  # pylint: disable=protected-access\n         for var in var_list:\n-            self.add_slot(var, \"slow\")\n+            self.add_slot(var, \"slow\", initializer=var)\n \n     def _create_hypers(self):\n         self._optimizer._create_hypers()  # pylint: disable=protected-access\n@@ -103,17 +103,6 @@\n         )  # pylint: disable=protected-access\n         return super().apply_gradients(grads_and_vars, name, **kwargs)\n \n-    def _init_op(self, var):\n-        slow_var = self.get_slot(var, \"slow\")\n-        return slow_var.assign(\n-            tf.where(\n-                tf.equal(self.iterations, tf.constant(0, dtype=self.iterations.dtype)),\n-                var,\n-                slow_var,\n-            ),\n-            use_locking=self._use_locking,\n-        )\n-\n     def _look_ahead_op(self, var):\n         var_dtype = var.dtype.base_dtype\n         slow_var = self.get_slot(var, \"slow\")\n@@ -138,24 +127,22 @@\n         return self._weights + self._optimizer.weights\n \n     def _resource_apply_dense(self, grad, var):\n-        init_op = self._init_op(var)\n-        with tf.control_dependencies([init_op]):\n-            train_op = self._optimizer._resource_apply_dense(\n-                grad, var\n-            )  # pylint: disable=protected-access\n-            with tf.control_dependencies([train_op]):\n-                look_ahead_op = self._look_ahead_op(var)\n-        return tf.group(init_op, train_op, look_ahead_op)\n+        train_op = self._optimizer._resource_apply_dense(\n+            grad, var\n+        )  # pylint: disable=protected-access\n+        with tf.control_dependencies([train_op]):\n+            look_ahead_op = self._look_ahead_op(var)\n+        return tf.group(train_op, look_ahead_op)\n \n     def _resource_apply_sparse(self, grad, var, indices):\n-        init_op = self._init_op(var)\n-        with tf.control_dependencies([init_op]):\n-            train_op = self._optimizer._resource_apply_sparse(  # pylint: disable=protected-access\n+        train_op = (\n+            self._optimizer._resource_apply_sparse(  # pylint: disable=protected-access\n                 grad, var, indices\n             )\n-            with tf.control_dependencies([train_op]):\n-                look_ahead_op = self._look_ahead_op(var)\n-        return tf.group(init_op, train_op, look_ahead_op)\n+        )\n+        with tf.control_dependencies([train_op]):\n+            look_ahead_op = self._look_ahead_op(var)\n+        return tf.group(train_op, look_ahead_op)\n \n     def get_config(self):\n         config = {\n", "issue": "Lookahead might divide weights with mixed_float16\n**System information**\r\n- OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Ubuntu 18.04\r\n- TensorFlow version and how it was installed (source or binary): 2.3.0-rc2 (source)\r\n- TensorFlow-Addons version and how it was installed (source or binary): Just lookahead extracted from the source (master)\r\n- Python version: 3.6.9\r\n- Is GPU used? (yes/no): yes\r\n\r\nWhen tf.keras.mixed_precision.experimental.set_policy('mixed_float16') is set, using LookAhead divides the layers weights after a few gradient updates (for example 7 in my case).\r\n\r\nAfter investigating, I think the reason is that when mixed_float16 is used with automatic loss scaling, the loss scale is initialized very high. Thus the first few gradient updates do nothing (see https://github.com/tensorflow/tensorflow/blob/582c8d236cb079023657287c318ff26adb239002/tensorflow/python/keras/mixed_precision/loss_scale_optimizer.py#L714 for how updates are skipped).\r\n\r\nAs a result, LookAhead ends up with its 'slow vars' initialized to 0 instead of the original layer weight. Indeed the vars are initialized to the correct value only on the first iteration during the gradient update (https://github.com/tensorflow/addons/blob/master/tensorflow_addons/optimizers/lookahead.py#L106).\r\n\r\nThus when the gradient update finally does something, the first variable update does average them with the slow variables (initialized to zero). Which would explain the weight division I'm seeing.\r\n\r\nI don't have a short code to replicate the bug, but I think the above explanation and links should highlight well enough the issue with the current way LookAhead's slow variables are initialized.\r\n\nLookahead might divide weights with mixed_float16\n**System information**\r\n- OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Ubuntu 18.04\r\n- TensorFlow version and how it was installed (source or binary): 2.3.0-rc2 (source)\r\n- TensorFlow-Addons version and how it was installed (source or binary): Just lookahead extracted from the source (master)\r\n- Python version: 3.6.9\r\n- Is GPU used? (yes/no): yes\r\n\r\nWhen tf.keras.mixed_precision.experimental.set_policy('mixed_float16') is set, using LookAhead divides the layers weights after a few gradient updates (for example 7 in my case).\r\n\r\nAfter investigating, I think the reason is that when mixed_float16 is used with automatic loss scaling, the loss scale is initialized very high. Thus the first few gradient updates do nothing (see https://github.com/tensorflow/tensorflow/blob/582c8d236cb079023657287c318ff26adb239002/tensorflow/python/keras/mixed_precision/loss_scale_optimizer.py#L714 for how updates are skipped).\r\n\r\nAs a result, LookAhead ends up with its 'slow vars' initialized to 0 instead of the original layer weight. Indeed the vars are initialized to the correct value only on the first iteration during the gradient update (https://github.com/tensorflow/addons/blob/master/tensorflow_addons/optimizers/lookahead.py#L106).\r\n\r\nThus when the gradient update finally does something, the first variable update does average them with the slow variables (initialized to zero). Which would explain the weight division I'm seeing.\r\n\r\nI don't have a short code to replicate the bug, but I think the above explanation and links should highlight well enough the issue with the current way LookAhead's slow variables are initialized.\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\nimport tensorflow as tf\nfrom tensorflow_addons.utils import types\n\nfrom typeguard import typechecked\n\n\n@tf.keras.utils.register_keras_serializable(package=\"Addons\")\nclass Lookahead(tf.keras.optimizers.Optimizer):\n    \"\"\"This class allows to extend optimizers with the lookahead mechanism.\n\n    The mechanism is proposed by Michael R. Zhang et.al in the paper\n    [Lookahead Optimizer: k steps forward, 1 step back]\n    (https://arxiv.org/abs/1907.08610v1). The optimizer iteratively updates two\n    sets of weights: the search directions for weights are chosen by the inner\n    optimizer, while the \"slow weights\" are updated each `k` steps based on the\n    directions of the \"fast weights\" and the two sets of weights are\n    synchronized. This method improves the learning stability and lowers the\n    variance of its inner optimizer.\n\n    Example of usage:\n\n    ```python\n    opt = tf.keras.optimizers.SGD(learning_rate)\n    opt = tfa.optimizers.Lookahead(opt)\n    ```\n    \"\"\"\n\n    @typechecked\n    def __init__(\n        self,\n        optimizer: types.Optimizer,\n        sync_period: int = 6,\n        slow_step_size: types.FloatTensorLike = 0.5,\n        name: str = \"Lookahead\",\n        **kwargs,\n    ):\n        r\"\"\"Wrap optimizer with the lookahead mechanism.\n\n        Args:\n            optimizer: The original optimizer that will be used to compute\n                and apply the gradients.\n            sync_period: An integer. The synchronization period of lookahead.\n                Enable lookahead mechanism by setting it with a positive value.\n            slow_step_size: A floating point value.\n                The ratio for updating the slow weights.\n            name: Optional name for the operations created when applying\n                gradients. Defaults to \"Lookahead\".\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\n        if isinstance(optimizer, str):\n            optimizer = tf.keras.optimizers.get(optimizer)\n        if not isinstance(optimizer, tf.keras.optimizers.Optimizer):\n            raise TypeError(\n                \"optimizer is not an object of tf.keras.optimizers.Optimizer\"\n            )\n\n        self._optimizer = optimizer\n        self._set_hyper(\"sync_period\", sync_period)\n        self._set_hyper(\"slow_step_size\", slow_step_size)\n        self._initialized = False\n        self._track_trackable(self._optimizer, \"lh_base_optimizer\")\n\n    def _create_slots(self, var_list):\n        self._optimizer._create_slots(\n            var_list=var_list\n        )  # pylint: disable=protected-access\n        for var in var_list:\n            self.add_slot(var, \"slow\")\n\n    def _create_hypers(self):\n        self._optimizer._create_hypers()  # pylint: disable=protected-access\n\n    def _prepare(self, var_list):\n        return self._optimizer._prepare(\n            var_list=var_list\n        )  # pylint: disable=protected-access\n\n    def apply_gradients(self, grads_and_vars, name=None, **kwargs):\n        self._optimizer._iterations = (\n            self.iterations\n        )  # pylint: disable=protected-access\n        return super().apply_gradients(grads_and_vars, name, **kwargs)\n\n    def _init_op(self, var):\n        slow_var = self.get_slot(var, \"slow\")\n        return slow_var.assign(\n            tf.where(\n                tf.equal(self.iterations, tf.constant(0, dtype=self.iterations.dtype)),\n                var,\n                slow_var,\n            ),\n            use_locking=self._use_locking,\n        )\n\n    def _look_ahead_op(self, var):\n        var_dtype = var.dtype.base_dtype\n        slow_var = self.get_slot(var, \"slow\")\n        local_step = tf.cast(self.iterations + 1, tf.dtypes.int64)\n        sync_period = self._get_hyper(\"sync_period\", tf.dtypes.int64)\n        slow_step_size = self._get_hyper(\"slow_step_size\", var_dtype)\n        step_back = slow_var + slow_step_size * (var - slow_var)\n        sync_cond = tf.equal(\n            tf.math.floordiv(local_step, sync_period) * sync_period, local_step\n        )\n        with tf.control_dependencies([step_back]):\n            slow_update = slow_var.assign(\n                tf.where(sync_cond, step_back, slow_var), use_locking=self._use_locking\n            )\n            var_update = var.assign(\n                tf.where(sync_cond, step_back, var), use_locking=self._use_locking\n            )\n        return tf.group(slow_update, var_update)\n\n    @property\n    def weights(self):\n        return self._weights + self._optimizer.weights\n\n    def _resource_apply_dense(self, grad, var):\n        init_op = self._init_op(var)\n        with tf.control_dependencies([init_op]):\n            train_op = self._optimizer._resource_apply_dense(\n                grad, var\n            )  # pylint: disable=protected-access\n            with tf.control_dependencies([train_op]):\n                look_ahead_op = self._look_ahead_op(var)\n        return tf.group(init_op, train_op, look_ahead_op)\n\n    def _resource_apply_sparse(self, grad, var, indices):\n        init_op = self._init_op(var)\n        with tf.control_dependencies([init_op]):\n            train_op = self._optimizer._resource_apply_sparse(  # pylint: disable=protected-access\n                grad, var, indices\n            )\n            with tf.control_dependencies([train_op]):\n                look_ahead_op = self._look_ahead_op(var)\n        return tf.group(init_op, train_op, look_ahead_op)\n\n    def get_config(self):\n        config = {\n            \"optimizer\": tf.keras.optimizers.serialize(self._optimizer),\n            \"sync_period\": self._serialize_hyperparameter(\"sync_period\"),\n            \"slow_step_size\": self._serialize_hyperparameter(\"slow_step_size\"),\n        }\n        base_config = super().get_config()\n        return {**base_config, **config}\n\n    @property\n    def learning_rate(self):\n        return self._optimizer._get_hyper(\"learning_rate\")\n\n    @learning_rate.setter\n    def learning_rate(self, learning_rate):\n        self._optimizer._set_hyper(\"learning_rate\", learning_rate)\n\n    @property\n    def lr(self):\n        return self.learning_rate\n\n    @lr.setter\n    def lr(self, lr):\n        self.learning_rate = lr\n\n    @classmethod\n    def from_config(cls, config, custom_objects=None):\n        optimizer = tf.keras.optimizers.deserialize(\n            config.pop(\"optimizer\"), custom_objects=custom_objects\n        )\n        return cls(optimizer, **config)\n", "path": "tensorflow_addons/optimizers/lookahead.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\nimport tensorflow as tf\nfrom tensorflow_addons.utils import types\n\nfrom typeguard import typechecked\n\n\n@tf.keras.utils.register_keras_serializable(package=\"Addons\")\nclass Lookahead(tf.keras.optimizers.Optimizer):\n    \"\"\"This class allows to extend optimizers with the lookahead mechanism.\n\n    The mechanism is proposed by Michael R. Zhang et.al in the paper\n    [Lookahead Optimizer: k steps forward, 1 step back]\n    (https://arxiv.org/abs/1907.08610v1). The optimizer iteratively updates two\n    sets of weights: the search directions for weights are chosen by the inner\n    optimizer, while the \"slow weights\" are updated each `k` steps based on the\n    directions of the \"fast weights\" and the two sets of weights are\n    synchronized. This method improves the learning stability and lowers the\n    variance of its inner optimizer.\n\n    Example of usage:\n\n    ```python\n    opt = tf.keras.optimizers.SGD(learning_rate)\n    opt = tfa.optimizers.Lookahead(opt)\n    ```\n    \"\"\"\n\n    @typechecked\n    def __init__(\n        self,\n        optimizer: types.Optimizer,\n        sync_period: int = 6,\n        slow_step_size: types.FloatTensorLike = 0.5,\n        name: str = \"Lookahead\",\n        **kwargs,\n    ):\n        r\"\"\"Wrap optimizer with the lookahead mechanism.\n\n        Args:\n            optimizer: The original optimizer that will be used to compute\n                and apply the gradients.\n            sync_period: An integer. The synchronization period of lookahead.\n                Enable lookahead mechanism by setting it with a positive value.\n            slow_step_size: A floating point value.\n                The ratio for updating the slow weights.\n            name: Optional name for the operations created when applying\n                gradients. Defaults to \"Lookahead\".\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\n        if isinstance(optimizer, str):\n            optimizer = tf.keras.optimizers.get(optimizer)\n        if not isinstance(optimizer, tf.keras.optimizers.Optimizer):\n            raise TypeError(\n                \"optimizer is not an object of tf.keras.optimizers.Optimizer\"\n            )\n\n        self._optimizer = optimizer\n        self._set_hyper(\"sync_period\", sync_period)\n        self._set_hyper(\"slow_step_size\", slow_step_size)\n        self._initialized = False\n        self._track_trackable(self._optimizer, \"lh_base_optimizer\")\n\n    def _create_slots(self, var_list):\n        self._optimizer._create_slots(\n            var_list=var_list\n        )  # pylint: disable=protected-access\n        for var in var_list:\n            self.add_slot(var, \"slow\", initializer=var)\n\n    def _create_hypers(self):\n        self._optimizer._create_hypers()  # pylint: disable=protected-access\n\n    def _prepare(self, var_list):\n        return self._optimizer._prepare(\n            var_list=var_list\n        )  # pylint: disable=protected-access\n\n    def apply_gradients(self, grads_and_vars, name=None, **kwargs):\n        self._optimizer._iterations = (\n            self.iterations\n        )  # pylint: disable=protected-access\n        return super().apply_gradients(grads_and_vars, name, **kwargs)\n\n    def _look_ahead_op(self, var):\n        var_dtype = var.dtype.base_dtype\n        slow_var = self.get_slot(var, \"slow\")\n        local_step = tf.cast(self.iterations + 1, tf.dtypes.int64)\n        sync_period = self._get_hyper(\"sync_period\", tf.dtypes.int64)\n        slow_step_size = self._get_hyper(\"slow_step_size\", var_dtype)\n        step_back = slow_var + slow_step_size * (var - slow_var)\n        sync_cond = tf.equal(\n            tf.math.floordiv(local_step, sync_period) * sync_period, local_step\n        )\n        with tf.control_dependencies([step_back]):\n            slow_update = slow_var.assign(\n                tf.where(sync_cond, step_back, slow_var), use_locking=self._use_locking\n            )\n            var_update = var.assign(\n                tf.where(sync_cond, step_back, var), use_locking=self._use_locking\n            )\n        return tf.group(slow_update, var_update)\n\n    @property\n    def weights(self):\n        return self._weights + self._optimizer.weights\n\n    def _resource_apply_dense(self, grad, var):\n        train_op = self._optimizer._resource_apply_dense(\n            grad, var\n        )  # pylint: disable=protected-access\n        with tf.control_dependencies([train_op]):\n            look_ahead_op = self._look_ahead_op(var)\n        return tf.group(train_op, look_ahead_op)\n\n    def _resource_apply_sparse(self, grad, var, indices):\n        train_op = (\n            self._optimizer._resource_apply_sparse(  # pylint: disable=protected-access\n                grad, var, indices\n            )\n        )\n        with tf.control_dependencies([train_op]):\n            look_ahead_op = self._look_ahead_op(var)\n        return tf.group(train_op, look_ahead_op)\n\n    def get_config(self):\n        config = {\n            \"optimizer\": tf.keras.optimizers.serialize(self._optimizer),\n            \"sync_period\": self._serialize_hyperparameter(\"sync_period\"),\n            \"slow_step_size\": self._serialize_hyperparameter(\"slow_step_size\"),\n        }\n        base_config = super().get_config()\n        return {**base_config, **config}\n\n    @property\n    def learning_rate(self):\n        return self._optimizer._get_hyper(\"learning_rate\")\n\n    @learning_rate.setter\n    def learning_rate(self, learning_rate):\n        self._optimizer._set_hyper(\"learning_rate\", learning_rate)\n\n    @property\n    def lr(self):\n        return self.learning_rate\n\n    @lr.setter\n    def lr(self, lr):\n        self.learning_rate = lr\n\n    @classmethod\n    def from_config(cls, config, custom_objects=None):\n        optimizer = tf.keras.optimizers.deserialize(\n            config.pop(\"optimizer\"), custom_objects=custom_objects\n        )\n        return cls(optimizer, **config)\n", "path": "tensorflow_addons/optimizers/lookahead.py"}]}