We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Allow passing in a tensor to tfa.optimizers.MovingAverage num_updates
Currently, `tfa.optimizers.MovingAverage` has an assert that requires `num_updates` to be of type int, [shown here](https://github.com/tensorflow/addons/blob/604a70de563f8797984c9c3f002aff70bef6c90b/tensorflow_addons/optimizers/moving_average.py#L81). This prevents me from passing in an integer tensor that changes with the global step, which is officially supported by `tf.train.ExponentialMovingAverage`.

Can this assert be updated to handle this use case?

--- 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/moving_average.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 from __future__ import absolute_import
17 from __future__ import division
18 from __future__ import print_function
19 
20 import tensorflow as tf
21 
22 
23 @tf.keras.utils.register_keras_serializable(package='Addons')
24 class MovingAverage(tf.keras.optimizers.Optimizer):
25     """Optimizer that computes a moving average of the variables.
26 
27     Empirically it has been found that using the moving average of the trained
28     parameters of a deep network is better than using its trained parameters
29     directly. This optimizer allows you to compute this moving average and swap
30     the variables at save time so that any code outside of the training loop
31     will use by default the average values instead of the original ones.
32 
33     Example of usage:
34 
35     ```python
36     opt = tf.keras.optimizers.SGD(learning_rate)
37     opt = tfa.optimizers.MovingAverage(opt)
38 
39     ```
40     """
41 
42     def __init__(self,
43                  optimizer,
44                  average_decay=0.1,
45                  num_updates=None,
46                  sequential_update=True,
47                  name="MovingAverage",
48                  **kwargs):
49         """Construct a new MovingAverage optimizer.
50 
51         Args:
52             optimizer: str or `tf.keras.optimizers.Optimizer` that will be
53                 used to compute and apply gradients.
54             average_decay: float. Decay to use to maintain the moving averages
55                 of trained variables. See `tf.train.ExponentialMovingAverage`
56                 for details.
57             num_updates: Optional count of the number of updates applied to
58                 variables. See `tf.train.ExponentialMovingAverage` for details.
59             sequential_update: Bool. If False, will compute the moving average
60                 at the same time as the model is updated, potentially doing
61                 benign data races. If True, will update the moving average
62                 after gradient updates.
63             name: Optional name for the operations created when applying
64                 gradients. Defaults to "MovingAverage".
65             **kwargs: keyword arguments. Allowed to be {`clipnorm`,
66                 `clipvalue`, `lr`, `decay`}. `clipnorm` is clip gradients by
67                 norm; `clipvalue` is clip gradients by value, `decay` is
68                 included for backward compatibility to allow time inverse
69                 decay of learning rate. `lr` is included for backward
70                 compatibility, recommended to use `learning_rate` instead.
71         """
72         super(MovingAverage, self).__init__(name, **kwargs)
73 
74         if isinstance(optimizer, str):
75             optimizer = tf.keras.optimizers.get(optimizer)
76 
77         if not isinstance(optimizer, tf.keras.optimizers.Optimizer):
78             raise TypeError(
79                 "optimizer is not an object of tf.keras.optimizers.Optimizer")
80 
81         if num_updates is not None and not isinstance(num_updates, int):
82             raise TypeError("num_updates must be None or of integer type")
83 
84         if not isinstance(sequential_update, bool):
85             raise TypeError("sequential_update must be of bool type")
86 
87         with tf.name_scope(name):
88             self._ema = tf.train.ExponentialMovingAverage(
89                 average_decay, num_updates=num_updates)
90 
91         self._optimizer = optimizer
92         self._set_hyper("average_decay", average_decay)
93         self._num_updates = num_updates
94         self._sequential_update = sequential_update
95         self._initialized = False
96 
97     def apply_gradients(self, grads_and_vars, name=None):
98         var_list = [v for (_, v) in grads_and_vars]
99 
100         if tf.executing_eagerly() and not self._initialized:
101             # this to ensure that var_list is registered initially
102             self._ema.apply(var_list)
103             self._initialized = True
104 
105         train_op = self._optimizer.apply_gradients(grads_and_vars, name=name)
106 
107         if self._sequential_update:
108             with tf.control_dependencies([train_op]):
109                 ma_op = self._ema.apply(var_list)
110         else:
111             ma_op = self._ema.apply(var_list)
112 
113         return tf.group(train_op, ma_op, name="train_with_avg")
114 
115     def get_config(self):
116         config = {
117             'optimizer': tf.keras.optimizers.serialize(self._optimizer),
118             'average_decay': self._serialize_hyperparameter('average_decay'),
119             'num_updates': self._num_updates,
120             'sequential_update': self._sequential_update
121         }
122         base_config = super(MovingAverage, self).get_config()
123         return dict(list(base_config.items()) + list(config.items()))
124 
125     @classmethod
126     def from_config(cls, config, custom_objects=None):
127         optimizer = tf.keras.optimizers.deserialize(
128             config.pop('optimizer'), custom_objects=custom_objects)
129         return cls(optimizer, **config)
130 
131     def assign_average_vars(self, var_list):
132         """Assign variables in var_list with their respective moving averages.
133 
134         Example:
135         ```python
136         model = tf.Sequential([...])
137         opt = tfa.optimizers.MovingAverage(
138             tf.keras.optimizers.SGD(lr=2.0), 0.5)
139 
140         model.compile(opt, ...)
141         model.fit(x, y, ...)
142 
143         # Update the weights to their mean before saving
144         opt.assign_average_vars(model.variables)
145 
146         model.save('model.h5')
147         ```
148         """
149         assign = tf.group([v.assign(self._ema.average(v)) for v in var_list])
150         return assign
151 
152     @property
153     def weights(self):
154         return self._optimizer.weights
155 
156     def _resource_apply_dense(self, grad, var):
157         return self._optimizer._resource_apply_dense(grad, var)  # pylint: disable=protected-access
158 
159     def _resource_apply_sparse_duplicate_indices(self, grad, var, indices):
160         return self._optimizer._resource_apply_sparse_duplicate_indices(  # pylint: disable=protected-access
161             grad, var, indices)
162 
163     def _resource_apply_sparse(self, grad, var, indices):
164         return self._optimizer._resource_apply_sparse(grad, var, indices)  # pylint: disable=protected-access
165 
166     @property
167     def learning_rate(self):
168         return self._optimizer._get_hyper('learning_rate')
169 
170     @learning_rate.setter
171     def learning_rate(self, learning_rate):
172         self._optimizer._set_hyper('learning_rate', learning_rate)
173 
174     @property
175     def lr(self):
176         return self.learning_rate
177 
178     @lr.setter
179     def lr(self, lr):
180         self.learning_rate = lr
181
```
--- 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/moving_average.py b/tensorflow_addons/optimizers/moving_average.py
--- a/tensorflow_addons/optimizers/moving_average.py
+++ b/tensorflow_addons/optimizers/moving_average.py
@@ -78,9 +78,6 @@
             raise TypeError(
                 "optimizer is not an object of tf.keras.optimizers.Optimizer")
 
-        if num_updates is not None and not isinstance(num_updates, int):
-            raise TypeError("num_updates must be None or of integer type")
-
         if not isinstance(sequential_update, bool):
             raise TypeError("sequential_update must be of bool type")

{"golden_diff": "diff --git a/tensorflow_addons/optimizers/moving_average.py b/tensorflow_addons/optimizers/moving_average.py\n--- a/tensorflow_addons/optimizers/moving_average.py\n+++ b/tensorflow_addons/optimizers/moving_average.py\n@@ -78,9 +78,6 @@\n             raise TypeError(\n                 \"optimizer is not an object of tf.keras.optimizers.Optimizer\")\n \n-        if num_updates is not None and not isinstance(num_updates, int):\n-            raise TypeError(\"num_updates must be None or of integer type\")\n-\n         if not isinstance(sequential_update, bool):\n             raise TypeError(\"sequential_update must be of bool type\")\n", "issue": "Allow passing in a tensor to tfa.optimizers.MovingAverage num_updates\nCurrently, `tfa.optimizers.MovingAverage` has an assert that requires `num_updates` to be of type int, [shown here](https://github.com/tensorflow/addons/blob/604a70de563f8797984c9c3f002aff70bef6c90b/tensorflow_addons/optimizers/moving_average.py#L81). This prevents me from passing in an integer tensor that changes with the global step, which is officially supported by `tf.train.ExponentialMovingAverage`.\r\n\r\nCan this assert be updated to handle this use case?\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\nfrom __future__ import absolute_import\nfrom __future__ import division\nfrom __future__ import print_function\n\nimport tensorflow as tf\n\n\n@tf.keras.utils.register_keras_serializable(package='Addons')\nclass MovingAverage(tf.keras.optimizers.Optimizer):\n    \"\"\"Optimizer that computes a moving average of the variables.\n\n    Empirically it has been found that using the moving average of the trained\n    parameters of a deep network is better than using its trained parameters\n    directly. This optimizer allows you to compute this moving average and swap\n    the variables at save time so that any code outside of the training loop\n    will use by default the average values instead of the original ones.\n\n    Example of usage:\n\n    ```python\n    opt = tf.keras.optimizers.SGD(learning_rate)\n    opt = tfa.optimizers.MovingAverage(opt)\n\n    ```\n    \"\"\"\n\n    def __init__(self,\n                 optimizer,\n                 average_decay=0.1,\n                 num_updates=None,\n                 sequential_update=True,\n                 name=\"MovingAverage\",\n                 **kwargs):\n        \"\"\"Construct a new MovingAverage optimizer.\n\n        Args:\n            optimizer: str or `tf.keras.optimizers.Optimizer` that will be\n                used to compute and apply gradients.\n            average_decay: float. Decay to use to maintain the moving averages\n                of trained variables. See `tf.train.ExponentialMovingAverage`\n                for details.\n            num_updates: Optional count of the number of updates applied to\n                variables. See `tf.train.ExponentialMovingAverage` for details.\n            sequential_update: Bool. If False, will compute the moving average\n                at the same time as the model is updated, potentially doing\n                benign data races. If True, will update the moving average\n                after gradient updates.\n            name: Optional name for the operations created when applying\n                gradients. Defaults to \"MovingAverage\".\n            **kwargs: keyword arguments. Allowed to be {`clipnorm`,\n                `clipvalue`, `lr`, `decay`}. `clipnorm` is clip gradients by\n                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(MovingAverage, self).__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        if num_updates is not None and not isinstance(num_updates, int):\n            raise TypeError(\"num_updates must be None or of integer type\")\n\n        if not isinstance(sequential_update, bool):\n            raise TypeError(\"sequential_update must be of bool type\")\n\n        with tf.name_scope(name):\n            self._ema = tf.train.ExponentialMovingAverage(\n                average_decay, num_updates=num_updates)\n\n        self._optimizer = optimizer\n        self._set_hyper(\"average_decay\", average_decay)\n        self._num_updates = num_updates\n        self._sequential_update = sequential_update\n        self._initialized = False\n\n    def apply_gradients(self, grads_and_vars, name=None):\n        var_list = [v for (_, v) in grads_and_vars]\n\n        if tf.executing_eagerly() and not self._initialized:\n            # this to ensure that var_list is registered initially\n            self._ema.apply(var_list)\n            self._initialized = True\n\n        train_op = self._optimizer.apply_gradients(grads_and_vars, name=name)\n\n        if self._sequential_update:\n            with tf.control_dependencies([train_op]):\n                ma_op = self._ema.apply(var_list)\n        else:\n            ma_op = self._ema.apply(var_list)\n\n        return tf.group(train_op, ma_op, name=\"train_with_avg\")\n\n    def get_config(self):\n        config = {\n            'optimizer': tf.keras.optimizers.serialize(self._optimizer),\n            'average_decay': self._serialize_hyperparameter('average_decay'),\n            'num_updates': self._num_updates,\n            'sequential_update': self._sequential_update\n        }\n        base_config = super(MovingAverage, self).get_config()\n        return dict(list(base_config.items()) + list(config.items()))\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        return cls(optimizer, **config)\n\n    def assign_average_vars(self, var_list):\n        \"\"\"Assign variables in var_list with their respective moving averages.\n\n        Example:\n        ```python\n        model = tf.Sequential([...])\n        opt = tfa.optimizers.MovingAverage(\n            tf.keras.optimizers.SGD(lr=2.0), 0.5)\n\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 = tf.group([v.assign(self._ema.average(v)) for v in var_list])\n        return assign\n\n    @property\n    def weights(self):\n        return self._optimizer.weights\n\n    def _resource_apply_dense(self, grad, var):\n        return self._optimizer._resource_apply_dense(grad, var)  # pylint: disable=protected-access\n\n    def _resource_apply_sparse_duplicate_indices(self, grad, var, indices):\n        return self._optimizer._resource_apply_sparse_duplicate_indices(  # pylint: disable=protected-access\n            grad, var, indices)\n\n    def _resource_apply_sparse(self, grad, var, indices):\n        return self._optimizer._resource_apply_sparse(grad, var, indices)  # pylint: disable=protected-access\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", "path": "tensorflow_addons/optimizers/moving_average.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\nfrom __future__ import absolute_import\nfrom __future__ import division\nfrom __future__ import print_function\n\nimport tensorflow as tf\n\n\n@tf.keras.utils.register_keras_serializable(package='Addons')\nclass MovingAverage(tf.keras.optimizers.Optimizer):\n    \"\"\"Optimizer that computes a moving average of the variables.\n\n    Empirically it has been found that using the moving average of the trained\n    parameters of a deep network is better than using its trained parameters\n    directly. This optimizer allows you to compute this moving average and swap\n    the variables at save time so that any code outside of the training loop\n    will use by default the average values instead of the original ones.\n\n    Example of usage:\n\n    ```python\n    opt = tf.keras.optimizers.SGD(learning_rate)\n    opt = tfa.optimizers.MovingAverage(opt)\n\n    ```\n    \"\"\"\n\n    def __init__(self,\n                 optimizer,\n                 average_decay=0.1,\n                 num_updates=None,\n                 sequential_update=True,\n                 name=\"MovingAverage\",\n                 **kwargs):\n        \"\"\"Construct a new MovingAverage optimizer.\n\n        Args:\n            optimizer: str or `tf.keras.optimizers.Optimizer` that will be\n                used to compute and apply gradients.\n            average_decay: float. Decay to use to maintain the moving averages\n                of trained variables. See `tf.train.ExponentialMovingAverage`\n                for details.\n            num_updates: Optional count of the number of updates applied to\n                variables. See `tf.train.ExponentialMovingAverage` for details.\n            sequential_update: Bool. If False, will compute the moving average\n                at the same time as the model is updated, potentially doing\n                benign data races. If True, will update the moving average\n                after gradient updates.\n            name: Optional name for the operations created when applying\n                gradients. Defaults to \"MovingAverage\".\n            **kwargs: keyword arguments. Allowed to be {`clipnorm`,\n                `clipvalue`, `lr`, `decay`}. `clipnorm` is clip gradients by\n                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(MovingAverage, self).__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        if not isinstance(sequential_update, bool):\n            raise TypeError(\"sequential_update must be of bool type\")\n\n        with tf.name_scope(name):\n            self._ema = tf.train.ExponentialMovingAverage(\n                average_decay, num_updates=num_updates)\n\n        self._optimizer = optimizer\n        self._set_hyper(\"average_decay\", average_decay)\n        self._num_updates = num_updates\n        self._sequential_update = sequential_update\n        self._initialized = False\n\n    def apply_gradients(self, grads_and_vars, name=None):\n        var_list = [v for (_, v) in grads_and_vars]\n\n        if tf.executing_eagerly() and not self._initialized:\n            # this to ensure that var_list is registered initially\n            self._ema.apply(var_list)\n            self._initialized = True\n\n        train_op = self._optimizer.apply_gradients(grads_and_vars, name=name)\n\n        if self._sequential_update:\n            with tf.control_dependencies([train_op]):\n                ma_op = self._ema.apply(var_list)\n        else:\n            ma_op = self._ema.apply(var_list)\n\n        return tf.group(train_op, ma_op, name=\"train_with_avg\")\n\n    def get_config(self):\n        config = {\n            'optimizer': tf.keras.optimizers.serialize(self._optimizer),\n            'average_decay': self._serialize_hyperparameter('average_decay'),\n            'num_updates': self._num_updates,\n            'sequential_update': self._sequential_update\n        }\n        base_config = super(MovingAverage, self).get_config()\n        return dict(list(base_config.items()) + list(config.items()))\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        return cls(optimizer, **config)\n\n    def assign_average_vars(self, var_list):\n        \"\"\"Assign variables in var_list with their respective moving averages.\n\n        Example:\n        ```python\n        model = tf.Sequential([...])\n        opt = tfa.optimizers.MovingAverage(\n            tf.keras.optimizers.SGD(lr=2.0), 0.5)\n\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 = tf.group([v.assign(self._ema.average(v)) for v in var_list])\n        return assign\n\n    @property\n    def weights(self):\n        return self._optimizer.weights\n\n    def _resource_apply_dense(self, grad, var):\n        return self._optimizer._resource_apply_dense(grad, var)  # pylint: disable=protected-access\n\n    def _resource_apply_sparse_duplicate_indices(self, grad, var, indices):\n        return self._optimizer._resource_apply_sparse_duplicate_indices(  # pylint: disable=protected-access\n            grad, var, indices)\n\n    def _resource_apply_sparse(self, grad, var, indices):\n        return self._optimizer._resource_apply_sparse(grad, var, indices)  # pylint: disable=protected-access\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", "path": "tensorflow_addons/optimizers/moving_average.py"}]}