We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
AttributeError: 'TrackableWeightHandler' object has no attribute 'trainable'
**System information**
- OS Platform and Distribution (e.g., Linux Ubuntu 16.04): 18.04
- TensorFlow version and how it was installed (source or binary): 2.3.0
- TensorFlow-Addons version and how it was installed (source or binary): 0.11.2
- Python version: 3.6.8
- Is GPU used? (yes/no): yes

**Describe the bug**

```TrackableWeightHandler``` doesn't have trainable and assign attributes but also to be appended to weights and trainable_weights. When I use ```AverageModelCheckpoint```, I meet the bug. I'm not sure it's a bug of addons or tensorflow.
```TrackableWeightHandler``` is added in ```TextVectorization```

A clear and concise description of what the bug is.
```
Traceback (most recent call last):
  File "/usr/lib/python3.6/runpy.py", line 193, in _run_module_as_main
    "__main__", mod_spec)
  File "/usr/lib/python3.6/runpy.py", line 85, in _run_code
    exec(code, run_globals)
  File "/root/.vscode-server/extensions/ms-python.python-2020.11.371526539/pythonFiles/lib/python/debugpy/__main__.py", line 45, in <module>
    cli.main()
  File "/root/.vscode-server/extensions/ms-python.python-2020.11.371526539/pythonFiles/lib/python/debugpy/../debugpy/server/cli.py", line 430, in main
    run()
  File "/root/.vscode-server/extensions/ms-python.python-2020.11.371526539/pythonFiles/lib/python/debugpy/../debugpy/server/cli.py", line 267, in run_file
    runpy.run_path(options.target, run_name=compat.force_str("__main__"))
  File "/usr/lib/python3.6/runpy.py", line 263, in run_path
    pkg_name=pkg_name, script_name=fname)
  File "/usr/lib/python3.6/runpy.py", line 96, in _run_module_code
    mod_name, mod_spec, pkg_name, script_name)
  File "/usr/lib/python3.6/runpy.py", line 85, in _run_code
    exec(code, run_globals)
  File "/usr/local/srv/automl/efficientdet/nlp.py", line 65, in <module>
    tfa.callbacks.AverageModelCheckpoint(update_weights=True, filepath=os.path.join(MODEL_DIR, 'ckpt_moving_average'))])
  File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/training.py", line 108, in _method_wrapper
    return method(self, *args, **kwargs)
  File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/training.py", line 1137, in fit
    callbacks.on_epoch_end(epoch, epoch_logs)
  File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/callbacks.py", line 412, in on_epoch_end
    callback.on_epoch_end(epoch, logs)
  File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/callbacks.py", line 1249, in on_epoch_end
    self._save_model(epoch=epoch, logs=logs)
  File "/usr/local/lib/python3.6/dist-packages/tensorflow_addons/callbacks/average_model_checkpoint.py", line 76, in _save_model
    self.model.optimizer.assign_average_vars(self.model.trainable_weights)
  File "/usr/local/lib/python3.6/dist-packages/tensorflow_addons/optimizers/average_wrapper.py", line 125, in assign_average_vars
    for var in var_list
  File "/usr/local/lib/python3.6/dist-packages/tensorflow_addons/optimizers/average_wrapper.py", line 126, in <listcomp>
    if var.trainable
AttributeError: 'TrackableWeightHandler' object has no attribute 'trainable'
```

**Code to reproduce the issue**
https://colab.research.google.com/drive/1OM7kAA6aOcNSTTXpd74kUnw-EeeM2uH5?usp=sharing

Provide a reproducible test case that is the bare minimum necessary to generate the problem.

**Other info / logs**

Change ```model.variables``` to ```model.trainable_weights``` could fix it temporarily.



--- 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/average_wrapper.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 abc
17 
18 import tensorflow as tf
19 from tensorflow_addons.utils import types
20 
21 from typeguard import typechecked
22 
23 
24 class AveragedOptimizerWrapper(tf.keras.optimizers.Optimizer, metaclass=abc.ABCMeta):
25     @typechecked
26     def __init__(
27         self, optimizer: types.Optimizer, name: str = "AverageOptimizer", **kwargs
28     ):
29         super().__init__(name, **kwargs)
30 
31         if isinstance(optimizer, str):
32             optimizer = tf.keras.optimizers.get(optimizer)
33 
34         if not isinstance(optimizer, tf.keras.optimizers.Optimizer):
35             raise TypeError(
36                 "optimizer is not an object of tf.keras.optimizers.Optimizer"
37             )
38 
39         self._optimizer = optimizer
40         self._track_trackable(self._optimizer, "awg_optimizer")
41 
42     def _create_slots(self, var_list):
43         self._optimizer._create_slots(var_list=var_list)
44         for var in var_list:
45             self.add_slot(var, "average")
46 
47     def _create_hypers(self):
48         self._optimizer._create_hypers()
49 
50     def _prepare(self, var_list):
51         return self._optimizer._prepare(var_list=var_list)
52 
53     def apply_gradients(self, grads_and_vars, name=None, **kwargs):
54         self._optimizer._iterations = self.iterations
55         return super().apply_gradients(grads_and_vars, name, **kwargs)
56 
57     @abc.abstractmethod
58     def average_op(self, var, average_var):
59         raise NotImplementedError
60 
61     def _apply_average_op(self, train_op, var):
62         average_var = self.get_slot(var, "average")
63         return self.average_op(var, average_var)
64 
65     def _resource_apply_dense(self, grad, var):
66         train_op = self._optimizer._resource_apply_dense(grad, var)
67         average_op = self._apply_average_op(train_op, var)
68         return tf.group(train_op, average_op)
69 
70     def _resource_apply_sparse(self, grad, var, indices):
71         train_op = self._optimizer._resource_apply_sparse(grad, var, indices)
72         average_op = self._apply_average_op(train_op, var)
73         return tf.group(train_op, average_op)
74 
75     def _resource_apply_sparse_duplicate_indices(self, grad, var, indices):
76         train_op = self._optimizer._resource_apply_sparse_duplicate_indices(
77             grad, var, indices
78         )
79         average_op = self._apply_average_op(train_op, var)
80         return tf.group(train_op, average_op)
81 
82     def assign_average_vars(self, var_list):
83         """Assign variables in var_list with their respective averages.
84 
85         Args:
86             var_list: List of model variables to be assigned to their average.
87 
88         Returns:
89             assign_op: The op corresponding to the assignment operation of
90             variables to their average.
91 
92         Example:
93         ```python
94         model = tf.Sequential([...])
95         opt = tfa.optimizers.SWA(
96                 tf.keras.optimizers.SGD(lr=2.0), 100, 10)
97         model.compile(opt, ...)
98         model.fit(x, y, ...)
99 
100         # Update the weights to their mean before saving
101         opt.assign_average_vars(model.variables)
102 
103         model.save('model.h5')
104         ```
105         """
106         assign_op = tf.group(
107             [
108                 var.assign(self.get_slot(var, "average"), use_locking=self._use_locking)
109                 for var in var_list
110                 if var.trainable
111             ]
112         )
113         return assign_op
114 
115     def get_config(self):
116         config = {
117             "optimizer": tf.keras.optimizers.serialize(self._optimizer),
118         }
119         base_config = super().get_config()
120         return {**base_config, **config}
121 
122     @classmethod
123     def from_config(cls, config, custom_objects=None):
124         optimizer = tf.keras.optimizers.deserialize(
125             config.pop("optimizer"), custom_objects=custom_objects
126         )
127         return cls(optimizer, **config)
128 
129     @property
130     def weights(self):
131         return self._weights + self._optimizer.weights
132 
133     @property
134     def lr(self):
135         return self._optimizer._get_hyper("learning_rate")
136 
137     @lr.setter
138     def lr(self, lr):
139         self._optimizer._set_hyper("learning_rate", lr)  #
140 
141     @property
142     def learning_rate(self):
143         return self._optimizer._get_hyper("learning_rate")
144 
145     @learning_rate.setter
146     def learning_rate(self, learning_rate):
147         self._optimizer._set_hyper("learning_rate", learning_rate)
148
```
--- 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/average_wrapper.py b/tensorflow_addons/optimizers/average_wrapper.py
--- a/tensorflow_addons/optimizers/average_wrapper.py
+++ b/tensorflow_addons/optimizers/average_wrapper.py
@@ -14,6 +14,7 @@
 # ==============================================================================
 
 import abc
+import warnings
 
 import tensorflow as tf
 from tensorflow_addons.utils import types
@@ -103,14 +104,17 @@
         model.save('model.h5')
         ```
         """
-        assign_op = tf.group(
-            [
-                var.assign(self.get_slot(var, "average"), use_locking=self._use_locking)
-                for var in var_list
-                if var.trainable
-            ]
-        )
-        return assign_op
+        assign_ops = []
+        for var in var_list:
+            try:
+                assign_ops.append(
+                    var.assign(
+                        self.get_slot(var, "average"), use_locking=self._use_locking
+                    )
+                )
+            except Exception as e:
+                warnings.warn("Unable to assign average slot to {} : {}".format(var, e))
+        return tf.group(assign_ops)
 
     def get_config(self):
         config = {

{"golden_diff": "diff --git a/tensorflow_addons/optimizers/average_wrapper.py b/tensorflow_addons/optimizers/average_wrapper.py\n--- a/tensorflow_addons/optimizers/average_wrapper.py\n+++ b/tensorflow_addons/optimizers/average_wrapper.py\n@@ -14,6 +14,7 @@\n # ==============================================================================\n \n import abc\n+import warnings\n \n import tensorflow as tf\n from tensorflow_addons.utils import types\n@@ -103,14 +104,17 @@\n         model.save('model.h5')\n         ```\n         \"\"\"\n-        assign_op = tf.group(\n-            [\n-                var.assign(self.get_slot(var, \"average\"), use_locking=self._use_locking)\n-                for var in var_list\n-                if var.trainable\n-            ]\n-        )\n-        return assign_op\n+        assign_ops = []\n+        for var in var_list:\n+            try:\n+                assign_ops.append(\n+                    var.assign(\n+                        self.get_slot(var, \"average\"), use_locking=self._use_locking\n+                    )\n+                )\n+            except Exception as e:\n+                warnings.warn(\"Unable to assign average slot to {} : {}\".format(var, e))\n+        return tf.group(assign_ops)\n \n     def get_config(self):\n         config = {\n", "issue": "AttributeError: 'TrackableWeightHandler' object has no attribute 'trainable'\n**System information**\r\n- OS Platform and Distribution (e.g., Linux Ubuntu 16.04): 18.04\r\n- TensorFlow version and how it was installed (source or binary): 2.3.0\r\n- TensorFlow-Addons version and how it was installed (source or binary): 0.11.2\r\n- Python version: 3.6.8\r\n- Is GPU used? (yes/no): yes\r\n\r\n**Describe the bug**\r\n\r\n```TrackableWeightHandler``` doesn't have trainable and assign attributes but also to be appended to weights and trainable_weights. When I use ```AverageModelCheckpoint```, I meet the bug. I'm not sure it's a bug of addons or tensorflow.\r\n```TrackableWeightHandler``` is added in ```TextVectorization```\r\n\r\nA clear and concise description of what the bug is.\r\n```\r\nTraceback (most recent call last):\r\n  File \"/usr/lib/python3.6/runpy.py\", line 193, in _run_module_as_main\r\n    \"__main__\", mod_spec)\r\n  File \"/usr/lib/python3.6/runpy.py\", line 85, in _run_code\r\n    exec(code, run_globals)\r\n  File \"/root/.vscode-server/extensions/ms-python.python-2020.11.371526539/pythonFiles/lib/python/debugpy/__main__.py\", line 45, in <module>\r\n    cli.main()\r\n  File \"/root/.vscode-server/extensions/ms-python.python-2020.11.371526539/pythonFiles/lib/python/debugpy/../debugpy/server/cli.py\", line 430, in main\r\n    run()\r\n  File \"/root/.vscode-server/extensions/ms-python.python-2020.11.371526539/pythonFiles/lib/python/debugpy/../debugpy/server/cli.py\", line 267, in run_file\r\n    runpy.run_path(options.target, run_name=compat.force_str(\"__main__\"))\r\n  File \"/usr/lib/python3.6/runpy.py\", line 263, in run_path\r\n    pkg_name=pkg_name, script_name=fname)\r\n  File \"/usr/lib/python3.6/runpy.py\", line 96, in _run_module_code\r\n    mod_name, mod_spec, pkg_name, script_name)\r\n  File \"/usr/lib/python3.6/runpy.py\", line 85, in _run_code\r\n    exec(code, run_globals)\r\n  File \"/usr/local/srv/automl/efficientdet/nlp.py\", line 65, in <module>\r\n    tfa.callbacks.AverageModelCheckpoint(update_weights=True, filepath=os.path.join(MODEL_DIR, 'ckpt_moving_average'))])\r\n  File \"/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/training.py\", line 108, in _method_wrapper\r\n    return method(self, *args, **kwargs)\r\n  File \"/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/training.py\", line 1137, in fit\r\n    callbacks.on_epoch_end(epoch, epoch_logs)\r\n  File \"/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/callbacks.py\", line 412, in on_epoch_end\r\n    callback.on_epoch_end(epoch, logs)\r\n  File \"/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/callbacks.py\", line 1249, in on_epoch_end\r\n    self._save_model(epoch=epoch, logs=logs)\r\n  File \"/usr/local/lib/python3.6/dist-packages/tensorflow_addons/callbacks/average_model_checkpoint.py\", line 76, in _save_model\r\n    self.model.optimizer.assign_average_vars(self.model.trainable_weights)\r\n  File \"/usr/local/lib/python3.6/dist-packages/tensorflow_addons/optimizers/average_wrapper.py\", line 125, in assign_average_vars\r\n    for var in var_list\r\n  File \"/usr/local/lib/python3.6/dist-packages/tensorflow_addons/optimizers/average_wrapper.py\", line 126, in <listcomp>\r\n    if var.trainable\r\nAttributeError: 'TrackableWeightHandler' object has no attribute 'trainable'\r\n```\r\n\r\n**Code to reproduce the issue**\r\nhttps://colab.research.google.com/drive/1OM7kAA6aOcNSTTXpd74kUnw-EeeM2uH5?usp=sharing\r\n\r\nProvide a reproducible test case that is the bare minimum necessary to generate the problem.\r\n\r\n**Other info / logs**\r\n\r\nChange ```model.variables``` to ```model.trainable_weights``` could fix it temporarily.\r\n\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 abc\n\nimport tensorflow as tf\nfrom tensorflow_addons.utils import types\n\nfrom typeguard import typechecked\n\n\nclass AveragedOptimizerWrapper(tf.keras.optimizers.Optimizer, metaclass=abc.ABCMeta):\n    @typechecked\n    def __init__(\n        self, optimizer: types.Optimizer, name: str = \"AverageOptimizer\", **kwargs\n    ):\n        super().__init__(name, **kwargs)\n\n        if isinstance(optimizer, str):\n            optimizer = tf.keras.optimizers.get(optimizer)\n\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._track_trackable(self._optimizer, \"awg_optimizer\")\n\n    def _create_slots(self, var_list):\n        self._optimizer._create_slots(var_list=var_list)\n        for var in var_list:\n            self.add_slot(var, \"average\")\n\n    def _create_hypers(self):\n        self._optimizer._create_hypers()\n\n    def _prepare(self, var_list):\n        return self._optimizer._prepare(var_list=var_list)\n\n    def apply_gradients(self, grads_and_vars, name=None, **kwargs):\n        self._optimizer._iterations = self.iterations\n        return super().apply_gradients(grads_and_vars, name, **kwargs)\n\n    @abc.abstractmethod\n    def average_op(self, var, average_var):\n        raise NotImplementedError\n\n    def _apply_average_op(self, train_op, var):\n        average_var = self.get_slot(var, \"average\")\n        return self.average_op(var, average_var)\n\n    def _resource_apply_dense(self, grad, var):\n        train_op = self._optimizer._resource_apply_dense(grad, var)\n        average_op = self._apply_average_op(train_op, var)\n        return tf.group(train_op, average_op)\n\n    def _resource_apply_sparse(self, grad, var, indices):\n        train_op = self._optimizer._resource_apply_sparse(grad, var, indices)\n        average_op = self._apply_average_op(train_op, var)\n        return tf.group(train_op, average_op)\n\n    def _resource_apply_sparse_duplicate_indices(self, grad, var, indices):\n        train_op = self._optimizer._resource_apply_sparse_duplicate_indices(\n            grad, var, indices\n        )\n        average_op = self._apply_average_op(train_op, var)\n        return tf.group(train_op, average_op)\n\n    def assign_average_vars(self, var_list):\n        \"\"\"Assign variables in var_list with their respective averages.\n\n        Args:\n            var_list: List of model variables to be assigned to their average.\n\n        Returns:\n            assign_op: The op corresponding to the assignment operation of\n            variables to their average.\n\n        Example:\n        ```python\n        model = tf.Sequential([...])\n        opt = tfa.optimizers.SWA(\n                tf.keras.optimizers.SGD(lr=2.0), 100, 10)\n        model.compile(opt, ...)\n        model.fit(x, y, ...)\n\n        # Update the weights to their mean before saving\n        opt.assign_average_vars(model.variables)\n\n        model.save('model.h5')\n        ```\n        \"\"\"\n        assign_op = tf.group(\n            [\n                var.assign(self.get_slot(var, \"average\"), use_locking=self._use_locking)\n                for var in var_list\n                if var.trainable\n            ]\n        )\n        return assign_op\n\n    def get_config(self):\n        config = {\n            \"optimizer\": tf.keras.optimizers.serialize(self._optimizer),\n        }\n        base_config = super().get_config()\n        return {**base_config, **config}\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\n    @property\n    def weights(self):\n        return self._weights + self._optimizer.weights\n\n    @property\n    def lr(self):\n        return self._optimizer._get_hyper(\"learning_rate\")\n\n    @lr.setter\n    def lr(self, lr):\n        self._optimizer._set_hyper(\"learning_rate\", lr)  #\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", "path": "tensorflow_addons/optimizers/average_wrapper.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 abc\nimport warnings\n\nimport tensorflow as tf\nfrom tensorflow_addons.utils import types\n\nfrom typeguard import typechecked\n\n\nclass AveragedOptimizerWrapper(tf.keras.optimizers.Optimizer, metaclass=abc.ABCMeta):\n    @typechecked\n    def __init__(\n        self, optimizer: types.Optimizer, name: str = \"AverageOptimizer\", **kwargs\n    ):\n        super().__init__(name, **kwargs)\n\n        if isinstance(optimizer, str):\n            optimizer = tf.keras.optimizers.get(optimizer)\n\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._track_trackable(self._optimizer, \"awg_optimizer\")\n\n    def _create_slots(self, var_list):\n        self._optimizer._create_slots(var_list=var_list)\n        for var in var_list:\n            self.add_slot(var, \"average\")\n\n    def _create_hypers(self):\n        self._optimizer._create_hypers()\n\n    def _prepare(self, var_list):\n        return self._optimizer._prepare(var_list=var_list)\n\n    def apply_gradients(self, grads_and_vars, name=None, **kwargs):\n        self._optimizer._iterations = self.iterations\n        return super().apply_gradients(grads_and_vars, name, **kwargs)\n\n    @abc.abstractmethod\n    def average_op(self, var, average_var):\n        raise NotImplementedError\n\n    def _apply_average_op(self, train_op, var):\n        average_var = self.get_slot(var, \"average\")\n        return self.average_op(var, average_var)\n\n    def _resource_apply_dense(self, grad, var):\n        train_op = self._optimizer._resource_apply_dense(grad, var)\n        average_op = self._apply_average_op(train_op, var)\n        return tf.group(train_op, average_op)\n\n    def _resource_apply_sparse(self, grad, var, indices):\n        train_op = self._optimizer._resource_apply_sparse(grad, var, indices)\n        average_op = self._apply_average_op(train_op, var)\n        return tf.group(train_op, average_op)\n\n    def _resource_apply_sparse_duplicate_indices(self, grad, var, indices):\n        train_op = self._optimizer._resource_apply_sparse_duplicate_indices(\n            grad, var, indices\n        )\n        average_op = self._apply_average_op(train_op, var)\n        return tf.group(train_op, average_op)\n\n    def assign_average_vars(self, var_list):\n        \"\"\"Assign variables in var_list with their respective averages.\n\n        Args:\n            var_list: List of model variables to be assigned to their average.\n\n        Returns:\n            assign_op: The op corresponding to the assignment operation of\n            variables to their average.\n\n        Example:\n        ```python\n        model = tf.Sequential([...])\n        opt = tfa.optimizers.SWA(\n                tf.keras.optimizers.SGD(lr=2.0), 100, 10)\n        model.compile(opt, ...)\n        model.fit(x, y, ...)\n\n        # Update the weights to their mean before saving\n        opt.assign_average_vars(model.variables)\n\n        model.save('model.h5')\n        ```\n        \"\"\"\n        assign_ops = []\n        for var in var_list:\n            try:\n                assign_ops.append(\n                    var.assign(\n                        self.get_slot(var, \"average\"), use_locking=self._use_locking\n                    )\n                )\n            except Exception as e:\n                warnings.warn(\"Unable to assign average slot to {} : {}\".format(var, e))\n        return tf.group(assign_ops)\n\n    def get_config(self):\n        config = {\n            \"optimizer\": tf.keras.optimizers.serialize(self._optimizer),\n        }\n        base_config = super().get_config()\n        return {**base_config, **config}\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\n    @property\n    def weights(self):\n        return self._weights + self._optimizer.weights\n\n    @property\n    def lr(self):\n        return self._optimizer._get_hyper(\"learning_rate\")\n\n    @lr.setter\n    def lr(self, lr):\n        self._optimizer._set_hyper(\"learning_rate\", lr)  #\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", "path": "tensorflow_addons/optimizers/average_wrapper.py"}]}