fix typing of num_updates in moving_average

tensorflow_addons/optimizers/moving_average.py

@@ -18,7 +18,7 @@
 from tensorflow_addons.optimizers import AveragedOptimizerWrapper
 from tensorflow_addons.utils import types
 
-from typing import Optional
+from typing import Union
 from typeguard import typechecked
 
 
@@ -47,7 +47,7 @@ def __init__(
         optimizer: types.Optimizer,
         sequential_update: bool = True,
         average_decay: types.FloatTensorLike = 0.99,
-        num_updates: Optional[str] = None,
+        num_updates: Union[None, int, tf.Variable] = None,
         start_step: int = 0,
         dynamic_decay: bool = False,
         name: str = "MovingAverage",
@@ -82,6 +82,14 @@ def __init__(
         super().__init__(optimizer, sequential_update, name, **kwargs)
         self._num_updates = num_updates
         if self._num_updates is not None:
+            if isinstance(self._num_updates, tf.Variable):
+                tf.debugging.assert_integer(
+                    self._num_updates,
+                    (
+                        'type of argument "num_updates" must be '
+                        "int; got {} instead".format(self._num_updates.dtype)
+                    ),
+                )
             num_updates = tf.cast(self._num_updates, tf.float32, name="num_updates")
             average_decay = tf.minimum(
                 average_decay, (1.0 + num_updates) / (10.0 + num_updates)

tensorflow_addons/optimizers/tests/moving_average_test.py

@@ -68,6 +68,19 @@ def test_opt_failure():
         MovingAverage(base_opt, 0.5)
 
 
+@pytest.mark.usefixtures("maybe_run_functions_eagerly")
+def test_num_updates_valid():
+    for num_updates in [1, tf.Variable(1)]:
+        MovingAverage("sgd", num_updates=num_updates)
+
+
+@pytest.mark.usefixtures("maybe_run_functions_eagerly")
+def test_num_updates_invalid():
+    for num_updates in [1.0, tf.Variable(1.0), "a"]:
+        with pytest.raises(TypeError):
+            MovingAverage("sgd", num_updates=num_updates)
+
+
 @pytest.mark.usefixtures("maybe_run_functions_eagerly")
 def test_model_weights_update():
     grad = tf.Variable([[0.1]])